From e5f1ceba4e52fcbfaf07f1e739daf8ac330fafff Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 13 Oct 2026 02:07:05 -0700 Subject: [PATCH 01/59] [None][feat] Add DynamicEPLB kernels and runtime integration Add per-layer dynamic expert-parallel load balancing for the MegaMoE CuTe DSL backend, selected with moe_config.load_balancer.mode: per_layer_dynamic. - HALO-Q plan scheduler and in-switch TMA expert copy kernels under cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb, with THOP ops. - PyTorch runtime: the per-layer dynamic load balancer, the rebalance slot scheduler and shared-slot arena, warmup handoff of the submission owner, and rank-aligned teardown in MPI and Ray worker shutdown. - DeepSeek-V4 shared experts fused into the MegaMoE FC12 kernel on SM107. - Vendor bangyus/cutedsl_megamoe main and the HALO-Q scheduler sources through 3rdparty/vendor_sources.lock.yaml. The MegaMoE directory is byte-identical to upstream, and the vendor patch carries only the adaptations that TensorRT-LLM main already had. Signed-off-by: jintaop --- .pre-commit-config.yaml | 2 +- 3rdparty/vendor_patches/cutedsl-megamoe.patch | 910 ++++++ 3rdparty/vendor_sources.lock.yaml | 33 + cpp/tensorrt_llm/kernels/moe/CMakeLists.txt | 5 +- .../loadBalance/dynamicEplb/CMakeLists.txt | 72 + .../dynamicEplb/moeRebalanceHaloQ.cu | 2093 ++++++++++++++ .../dynamicEplb/moeRebalanceHaloQ.h | 67 + .../dynamicEplb/moeRebalancePlaceholder.cpp | 100 + .../dynamicEplb/moeRebalancePlanPolicy.h | 380 +++ .../dynamicEplb/moeRebalanceTma.cu | 1085 +++++++ .../loadBalance/dynamicEplb/moeRebalanceTma.h | 203 ++ .../dynamicEplb/moeRebalanceTmaGpuPlan.h | 53 + .../moeRebalanceTmaGpuPlanBuilder.cuh | 492 ++++ cpp/tensorrt_llm/thop/fp4Quantize.cpp | 22 +- cpp/tensorrt_llm/thop/fp8Quantize.cpp | 19 +- .../moe/loadBalance/moeRebalanceHaloQOp.cpp | 162 ++ .../moe/loadBalance/moeRebalanceTmaOp.cpp | 454 +++ pyproject.toml | 3 + .../ops/moe/__init__.py | 27 +- .../kernels/rubin/dsv4_qb_fusion/__init__.py | 24 +- .../_torch/custom_ops/cpp_custom_ops.py | 60 +- tensorrt_llm/_torch/model_config.py | 7 +- .../_torch/models/modeling_deepseekv4.py | 169 +- .../_torch/modules/megamoe_shared_mlp.py | 243 ++ .../custom_ops/cute_dsl_megamoe_custom_op.py | 621 +++- .../moe/fused_moe/MOE_DEVELOPER_GUIDE.md | 36 +- tensorrt_llm/_torch/moe/fused_moe/__init__.py | 4 + .../_torch/moe/fused_moe/create_moe.py | 7 + .../_torch/moe/fused_moe/impl_contract.py | 54 +- .../_torch/moe/fused_moe/interface.py | 84 +- .../fused_moe/mega_moe/mega_moe_cute_dsl.py | 418 ++- .../fused_moe/mega_moe/mega_moe_deepgemm.py | 24 +- .../mega_moe/rebalance_live_arena.py | 1058 +++++++ .../mega_moe/rebalance_slot_scheduler.py | 503 ++++ .../_torch/moe/fused_moe/moe_load_balancer.py | 86 +- .../_torch/moe/fused_moe/moe_resolution.py | 23 + .../_torch/moe/fused_moe/moe_scheduler.py | 89 +- .../moe/fused_moe/per_layer_dynamic_eplb.py | 564 ++++ .../_torch/moe/fused_moe/quantization.py | 284 +- .../moe/kernels/cutedsl_megamoe/VENDOR_STAMP | 7 + .../moe/kernels/cutedsl_megamoe/__init__.py | 25 +- .../_torch/moe/kernels/cutedsl_megamoe/api.py | 48 +- .../cutedsl_megamoe/communication/__init__.py | 49 +- .../communication/nvlink_domain/__init__.py | 28 +- .../nvlink_domain/symmetric_buffer.py | 60 +- .../communication/nvlink_domain/token_comm.py | 1601 ++++++---- .../nvlink_domain/token_comm_cta_mainloop.py | 93 + .../nvlink_domain/token_comm_deterministic.py | 2576 +++++++++++++++++ .../communication/token_protocol.py | 13 +- .../cutedsl_megamoe/helpers/__init__.py | 24 +- .../cutedsl_megamoe/helpers/constants.py | 4 +- .../helpers/cute_py_helpers.py | 46 +- .../helpers/device_workspace.py | 33 +- .../helpers/dlb_expert_count.py | 30 + .../cutedsl_megamoe/helpers/dsl_helpers.py | 66 +- .../helpers/fc2_ready_geometry.py | 36 + .../cutedsl_megamoe/helpers/flag_batch.py | 87 +- .../helpers/helper_weight_ready_gate.py | 318 ++ .../cutedsl_megamoe/helpers/iket_compat.py | 4 - .../cutedsl_megamoe/helpers/megamoe_aot.py | 14 + .../cutedsl_megamoe/helpers/ptx_helpers.py | 232 +- .../cutedsl_megamoe/helpers/smem_workspace.py | 76 +- .../cutedsl_megamoe/helpers/software_sync.py | 23 +- .../inference/mega => helpers}/topk_reduce.py | 151 +- .../kernels/cutedsl_megamoe/helpers/utils.py | 8 +- .../cutedsl_megamoe/kernel_src/__init__.py | 24 +- .../kernel_src/blackwell/__init__.py | 32 +- .../blackwell/custom_mix_cga_helpers.py | 105 +- .../blackwell/inference/__init__.py | 3 +- .../blackwell/inference/mega/__init__.py | 32 +- .../block_scaled_swap_ab_fc12_epilogue.py | 653 ++--- .../block_scaled_swap_ab_fc12_extension.py | 26 +- .../mega/block_scaled_swap_ab_fc12_kernel.py | 689 +++++ .../block_scaled_swap_ab_fc12_mainloop.py | 310 +- .../block_scaled_swap_ab_mega_moe_kernel.py | 324 +-- .../inference/mega/dynamic_mainloop.py | 24 +- .../kernel_src/function_mapping.py | 23 +- .../kernel_src/rubin/__init__.py | 8 +- .../rubin/custom_mix_cga_helpers.py | 25 + .../kernel_src/rubin/inference/__init__.py | 27 +- .../rubin/inference/local_mega/__init__.py | 33 +- .../block_scaled_swap_ab_fc12_epilogue.py | 96 +- .../block_scaled_swap_ab_fc12_extension.py | 7 + .../block_scaled_swap_ab_fc12_mainloop.py | 309 +- ...ck_scaled_swap_ab_local_mega_moe_kernel.py | 284 +- .../inference/local_mega/local_routing.py | 465 +-- .../rubin/inference/local_mega/tma_gather.py | 4 +- .../block_scaled_swap_ab_fc12_epilogue.py | 613 ++-- .../block_scaled_swap_ab_fc12_extension.py | 26 +- .../block_scaled_swap_ab_fc12_mainloop.py | 381 +-- ...d_swap_ab_fc12_mainloop_gen_specialized.py | 443 +-- .../block_scaled_swap_ab_mega_moe_kernel.py | 857 ++++-- ...swap_ab_mega_moe_kernel_gen_specialized.py | 299 +- .../rubin/inference/mega/dynamic_mainloop.py | 18 +- .../rubin/inference/mega/tma_gather.py | 7 + .../mega/token_comm_gen_specialized.py | 443 +-- .../kernel_src/schedulers/base.py | 32 +- .../kernel_src/schedulers/clc_mixed_cga.py | 162 ++ .../kernel_src/schedulers/fc12_mapping.py | 265 +- .../kernel_src/schedulers/fc12_scheduler.py | 454 ++- .../schedulers/non_clc_mixed_cga.py | 167 +- .../kernel_src/schedulers/work_id_claim.py | 444 ++- .../moe/kernels/cutedsl_megamoe/quant_def.py | 39 +- .../_torch/moe/kernels/megamoe_kernels.py | 23 + .../kernels/megamoe_scheduler/VENDOR_STAMP | 6 + .../moe/kernels/megamoe_scheduler/__init__.py | 46 + .../cuda_scheduler/__init__.py | 23 + .../cuda_scheduler/multirank_dist.py | 125 + .../cuda_scheduler/runtime.py | 607 ++++ .../moe/kernels/megamoe_scheduler/geometry.py | 35 + .../megamoe_scheduler/sami/__init__.py | 40 + .../kernels/megamoe_scheduler/sami/_util.py | 66 + .../kernels/megamoe_scheduler/sami/arena.py | 419 +++ .../kernels/megamoe_scheduler/sami/fabric.py | 900 ++++++ .../megamoe_scheduler/sami/geometry.py | 304 ++ .../megamoe_scheduler/sami/hierarchical.py | 695 +++++ .../_torch/moe/kernels/megamoe_shared_fc12.py | 652 +++++ .../_torch/pyexecutor/model_engine.py | 122 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 64 +- .../_torch/weight_sharing/source_identity.py | 6 +- tensorrt_llm/executor/base_worker.py | 63 +- tensorrt_llm/executor/ray/executor.py | 40 +- tensorrt_llm/executor/ray/gpu_worker.py | 122 +- tensorrt_llm/executor/rpc/rpc_server.py | 14 + tensorrt_llm/executor/rpc_proxy.py | 60 +- tensorrt_llm/executor/rpc_worker.py | 87 +- tensorrt_llm/executor/worker.py | 204 +- tensorrt_llm/llmapi/_load_balance_env.py | 49 + tensorrt_llm/llmapi/llm.py | 16 +- tensorrt_llm/llmapi/llm_args.py | 228 +- tensorrt_llm/llmapi/mpi_session.py | 9 +- tensorrt_llm/llmapi/serialization.py | 3 +- .../usage/llm_args_golden_manifest.json | 24 + .../integration/test_lists/test-db/l0_cpu.yml | 3 + .../test_megamoe_shared_fc12_autotune.py | 458 +++ .../modeling/test_modeling_deepseekv4.py | 119 + .../parallel/test_quantization_sm_budget.py | 110 + .../_torch/moe/test_dynamic_eplb_contracts.py | 1765 +++++++++++ .../_torch/moe/test_kimi_k3_situ_moe.py | 19 + .../_torch/moe/test_megamoe_streaming_load.py | 111 +- tests/unittest/_torch/moe/test_moe_backend.py | 363 ++- tests/unittest/_torch/moe/test_moe_impl.py | 29 + .../_torch/moe/test_per_layer_dynamic_eplb.py | 411 +++ .../weight_sharing/_source_identity_fakes.py | 2 + .../weight_sharing/test_source_identity.py | 16 + .../executor/test_shutdown_retry_contracts.py | 584 ++++ tests/unittest/llmapi/test_gms_args.py | 45 + tests/unittest/llmapi/test_llm_args.py | 202 +- tests/unittest/llmapi/test_serialization.py | 27 + .../usage/test_llmapi_config_capture.py | 15 +- 150 files changed, 27751 insertions(+), 5925 deletions(-) create mode 100644 3rdparty/vendor_patches/cutedsl-megamoe.patch create mode 100644 cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/CMakeLists.txt create mode 100644 cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.cu create mode 100644 cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h create mode 100644 cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlaceholder.cpp create mode 100644 cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlanPolicy.h create mode 100644 cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu create mode 100644 cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h create mode 100644 cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlan.h create mode 100644 cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlanBuilder.cuh create mode 100644 cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.cpp create mode 100644 cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp create mode 100644 tensorrt_llm/_torch/modules/megamoe_shared_mlp.py create mode 100644 tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_live_arena.py create mode 100644 tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py create mode 100644 tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py create mode 100644 tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/VENDOR_STAMP create mode 100644 tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_cta_mainloop.py create mode 100644 tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_deterministic.py create mode 100644 tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/dlb_expert_count.py create mode 100644 tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/fc2_ready_geometry.py create mode 100644 tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/helper_weight_ready_gate.py create mode 100644 tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/megamoe_aot.py rename tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/{kernel_src/blackwell/inference/mega => helpers}/topk_reduce.py (82%) create mode 100644 tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_kernel.py create mode 100644 tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/custom_mix_cga_helpers.py create mode 100644 tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_extension.py create mode 100644 tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/tma_gather.py create mode 100644 tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/clc_mixed_cga.py create mode 100644 tensorrt_llm/_torch/moe/kernels/megamoe_kernels.py create mode 100644 tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/VENDOR_STAMP create mode 100644 tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/__init__.py create mode 100644 tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/__init__.py create mode 100644 tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/multirank_dist.py create mode 100644 tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py create mode 100644 tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/geometry.py create mode 100644 tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/__init__.py create mode 100644 tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/_util.py create mode 100644 tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/arena.py create mode 100644 tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/fabric.py create mode 100644 tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/geometry.py create mode 100644 tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py create mode 100644 tensorrt_llm/_torch/moe/kernels/megamoe_shared_fc12.py create mode 100644 tensorrt_llm/llmapi/_load_balance_env.py create mode 100644 tests/unittest/_torch/cute_dsl_kernels/test_megamoe_shared_fc12_autotune.py create mode 100644 tests/unittest/_torch/moe/kernels/parallel/test_quantization_sm_budget.py create mode 100644 tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py create mode 100644 tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py create mode 100644 tests/unittest/executor/test_shutdown_retry_contracts.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c8450f647eb0..b2e83942a02d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1463,7 +1463,7 @@ legacy-files: &legacy_files | # list; the hook's own `files:` pattern only gates *when* the hook triggers. # Global exclude: vendored code + trtllm-gen FMHA artifacts (cubin pointers, export headers, cuda_ptx) -exclude: '(^tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_chunked_bt16(_small_bh)?\.py$|^cpp/tensorrt_llm/common/sha256/|^triton_kernels/|^tensorrt_llm/_torch/attention/backends/prims_ts/|^tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/|trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo\.h$|trtllmGenKernels/gemm/trtllmGen_gemm_export/KernelMetaInfo\.h$|\.cubin\.tar\.zst$)' +exclude: '(^tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_chunked_bt16(_small_bh)?\.py$|^cpp/tensorrt_llm/common/sha256/|^triton_kernels/|^tensorrt_llm/_torch/attention/backends/prims_ts/|^tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/|^tensorrt_llm/_torch/cute_dsl_kernels/cutedsl_megamoe/|^tensorrt_llm/_torch/cute_dsl_kernels/megamoe_scheduler/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/|trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo\.h$|trtllmGenKernels/gemm/trtllmGen_gemm_export/KernelMetaInfo\.h$|\.cubin\.tar\.zst$)' default_install_hook_types: [pre-commit, commit-msg] repos: diff --git a/3rdparty/vendor_patches/cutedsl-megamoe.patch b/3rdparty/vendor_patches/cutedsl-megamoe.patch new file mode 100644 index 000000000000..cfc9eb9edb10 --- /dev/null +++ b/3rdparty/vendor_patches/cutedsl-megamoe.patch @@ -0,0 +1,910 @@ +diff --git a/communication/nvlink_domain/token_comm.py b/communication/nvlink_domain/token_comm.py +index 7e2d5b2a70509a1a960fa450f95642ef321c9913..910e753adb0fcff4969ed89903d4d250cab06850 100644 +--- a/communication/nvlink_domain/token_comm.py ++++ b/communication/nvlink_domain/token_comm.py +@@ -113,7 +113,7 @@ class _MetadataPushRouter(KernelComponent): + sorted_metadata_region = "nvlink.token_comm.sorted_metadata" + sorted_scores_region = "nvlink.token_comm.sorted_scores" + pool_expert_base_region = "nvlink.token_comm.pool_expert_base" +- token_src_metadata_region = "nvlink.token_comm.token_src_metadata" ++ token_src_metadata_region = "nvlink.token_comm.token_src_metadata" # nosec B105 + fc1_topk_scores_region = "nvlink.token_comm.fc1_topk_scores" + source_expert_base_region = "nvlink.token_comm.source_expert_base" + push_destination_base_region = "nvlink.token_comm.push_destination_base" +@@ -1087,16 +1087,16 @@ class TokenCommNonDeterministic(KernelComponent): + fc2_activation_sf_region = "nvlink.token_comm.fc2_activation_sf" + pre_reduced_activation_region = "nvlink.token_comm.pre_reduced_activation" + pre_reduced_activation_sf_region = "nvlink.token_comm.pre_reduced_activation_sf" +- token_back_schedule_region = "nvlink.token_comm.token_back_schedule" ++ token_back_schedule_region = "nvlink.token_comm.token_back_schedule" # nosec B105 + +- token_in_mbarrier_region = "nvlink.token_comm.main_smem.token_in_mbarriers" +- token_back_mbarrier_region = "nvlink.token_comm.main_smem.token_back_mbarriers" ++ token_in_mbarrier_region = "nvlink.token_comm.main_smem.token_in_mbarriers" # nosec B105 ++ token_back_mbarrier_region = "nvlink.token_comm.main_smem.token_back_mbarriers" # nosec B105 + expert_sizes_smem_region = "nvlink.token_comm.main_smem.expert_sizes" +- token_in_activation_smem_region = "nvlink.token_comm.main_smem.token_in_activation" +- token_in_sf_smem_region = "nvlink.token_comm.main_smem.token_in_sf" +- token_in_sf_group_region = "nvlink.token_comm.main_smem.token_in_sf_group" +- token_back_activation_smem_region = "nvlink.token_comm.main_smem.token_back_activation" +- token_back_sf_smem_region = "nvlink.token_comm.main_smem.token_back_sf" ++ token_in_activation_smem_region = "nvlink.token_comm.main_smem.token_in_activation" # nosec B105 ++ token_in_sf_smem_region = "nvlink.token_comm.main_smem.token_in_sf" # nosec B105 ++ token_in_sf_group_region = "nvlink.token_comm.main_smem.token_in_sf_group" # nosec B105 ++ token_back_activation_smem_region = "nvlink.token_comm.main_smem.token_back_activation" # nosec B105 ++ token_back_sf_smem_region = "nvlink.token_comm.main_smem.token_back_sf" # nosec B105 + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: +@@ -1184,7 +1184,7 @@ class TokenCommNonDeterministic(KernelComponent): + # FC12 occupies warps 0..7. Legacy primitive callers retain their + # modulo-four numbering by omitting the explicit CTA ordinal. + self.dispatch_warp_start = 8 +- self.combine_warp_start = 12 if self.token_back_mode == "standalone_warps" else 8 ++ self.combine_warp_start = 12 if self.token_back_mode == "standalone_warps" else 8 # nosec B105 + self.dispatch_warp_count = self.transfer_warp_count + self.combine_warp_count = self.transfer_warp_count + self.dispatch_cta_count = self.promised_launchable_sm_count +@@ -1261,9 +1261,9 @@ class TokenCommNonDeterministic(KernelComponent): + ) + if self.token_back_ready_granularity not in ("expert", "token_tile"): + raise ValueError(f"Unsupported token_back_ready_granularity {self.token_back_ready_granularity!r}.") +- if self.token_back_ready_granularity == "token_tile" and ( ++ if self.token_back_ready_granularity == "token_tile" and ( # nosec B105 + self.token_back_mode not in ("standalone_warps", "reuse_dispatch_warps") +- or self.token_back_schedule_mode != "atomic_counter" ++ or self.token_back_schedule_mode != "atomic_counter" # nosec B105 + ): + raise ValueError("token_tile return-ready requires standalone_warps or reuse_dispatch_warps with atomic_counter.") + if not 1 <= self.token_in_flag_batch <= 32: +@@ -1281,7 +1281,7 @@ class TokenCommNonDeterministic(KernelComponent): + raise ValueError("sf_padding_block must be a multiple of 128.") + + if self.token_back_enabled and self.combine_stage_depth > 1: +- if self.token_back_schedule_mode != "static": ++ if self.token_back_schedule_mode != "static": # nosec B105 + raise ValueError("Multi-stage token-back requires static scheduling.") + if self.combine_format.is_quantized: + raise ValueError("Multi-stage token-back supports BF16 combine only.") +@@ -1322,7 +1322,7 @@ class TokenCommNonDeterministic(KernelComponent): + + @property + def token_back_push_data(self) -> bool: +- return self.token_back_mode != "epi_warps" ++ return self.token_back_mode != "epi_warps" # nosec B105 + + @property + def token_back_push_sf(self) -> bool: +@@ -1540,7 +1540,7 @@ class TokenCommNonDeterministic(KernelComponent): + if self.token_back_enabled: + fc2_done_slots = ( + self.max_fc1_ready_slot_count +- if self.token_back_ready_granularity == "token_tile" ++ if self.token_back_ready_granularity == "token_tile" # nosec B105 + else self.experts_per_rank + ) + workspace.register( +@@ -1564,7 +1564,7 @@ class TokenCommNonDeterministic(KernelComponent): + buffer_space="local", + byte_alignment=128, + ) +- if self.token_back_enabled and self.token_back_schedule_mode == "atomic_counter": ++ if self.token_back_enabled and self.token_back_schedule_mode == "atomic_counter": # nosec B105 + workspace.register( + self.token_back_schedule_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) +@@ -1609,7 +1609,7 @@ class TokenCommNonDeterministic(KernelComponent): + self._register_token_in_staging(transfer_overlay.add_lifetime("token_in")) + if not self.token_back_enabled: + return +- if self.token_back_mode == "standalone_warps": ++ if self.token_back_mode == "standalone_warps": # nosec B105 + back_overlay = workspace.create_overlay("nvlink.token_comm.main_smem.standalone_token_back") + else: + back_overlay = transfer_overlay +@@ -1664,20 +1664,20 @@ class TokenCommNonDeterministic(KernelComponent): + combine_rings = self.combine_warp_count + combine_depth = self.combine_stage_depth + sf_padded = self.activation_sf_hidden_padded +- if self.token_back_mode == "standalone_warps": ++ if self.token_back_mode == "standalone_warps": # nosec B105 + available_bytes_per_warp = self.token_back_chunk_bytes + else: + available_bytes_per_warp = self.bytes_per_token + sf_padded * int(self.activation_sf_dtype.width) // 8 + + if self.token_back_push_data: + bytes_per_output_token = self.hidden_size * int(self.combine_format.act_dtype.width) // 8 +- if self.token_back_mode == "standalone_warps": ++ if self.token_back_mode == "standalone_warps": # nosec B105 + chunk_bytes = self.token_back_chunk_bytes + elif available_bytes_per_warp < bytes_per_output_token: + chunk_bytes = self.bytes_per_token + else: + chunk_bytes = bytes_per_output_token +- if self.token_back_mode != "standalone_warps" and bytes_per_output_token % chunk_bytes != 0: ++ if self.token_back_mode != "standalone_warps" and bytes_per_output_token % chunk_bytes != 0: # nosec B105 + raise ValueError("Token-back data chunk bytes must divide one row.") + chunk_elements = chunk_bytes * 8 // int(self.combine_format.act_dtype.width) + shape = ( +@@ -2176,7 +2176,7 @@ class TokenCommNonDeterministic(KernelComponent): + ) + token_in_size_barrier.arrive_and_wait() + iket.range_pop() +- if cutlass.const_expr(not self.cta_specialized and self.token_back_mode == "standalone_warps"): ++ if cutlass.const_expr(not self.cta_specialized and self.token_back_mode == "standalone_warps"): # nosec B105 + sizes_ready_barrier = pipeline.NamedBarrier( + barrier_id=self.standalone_size_barrier_id, + num_threads=2 * self.transfer_thread_count, +@@ -2335,7 +2335,7 @@ class TokenCommNonDeterministic(KernelComponent): + + flag_tracker.fire() + cute.arch.sync_warp() +- if cutlass.const_expr(self.token_back_enabled and self.token_back_mode != "standalone_warps"): ++ if cutlass.const_expr(self.token_back_enabled and self.token_back_mode != "standalone_warps"): # nosec B105 + iket.range_push("token_in.transfer_barrier") + transfer_lifetime_barrier = pipeline.NamedBarrier( + barrier_id=self.transfer_lifetime_barrier_id, +@@ -2461,7 +2461,7 @@ class TokenCommNonDeterministic(KernelComponent): + cute.arch.mbarrier_init(worker_mbarrier + Int32(init_stage), 1) + cute.arch.sync_warp() + +- if cutlass.const_expr(not self.cta_specialized and self.token_back_mode == "standalone_warps"): ++ if cutlass.const_expr(not self.cta_specialized and self.token_back_mode == "standalone_warps"): # nosec B105 + sizes_ready_barrier = pipeline.NamedBarrier( + barrier_id=self.standalone_size_barrier_id, + num_threads=2 * self.transfer_thread_count, +@@ -2659,7 +2659,7 @@ class TokenCommNonDeterministic(KernelComponent): + cute.arch.mbarrier_init(worker_mbarrier, 1) + cute.arch.sync_warp() + +- if cutlass.const_expr(not self.cta_specialized and self.token_back_mode == "standalone_warps"): ++ if cutlass.const_expr(not self.cta_specialized and self.token_back_mode == "standalone_warps"): # nosec B105 + sizes_ready_barrier = pipeline.NamedBarrier( + barrier_id=self.standalone_size_barrier_id, + num_threads=2 * self.transfer_thread_count, +@@ -2675,7 +2675,7 @@ class TokenCommNonDeterministic(KernelComponent): + activation_chunk_count = ceil_div(output_token_bytes, activation_chunk_bytes) + data_window_unit = activation_chunk_bytes * 2 + reuse_data_pacing_enabled = ( +- self.token_back_mode == "reuse_dispatch_warps" and data_window_unit > self.minimum_pacing_window_cycles ++ self.token_back_mode == "reuse_dispatch_warps" and data_window_unit > self.minimum_pacing_window_cycles # nosec B105 + ) + # Preserve the empirical low:initial:high ratio of 1:2.5:5. + data_average_window = Int32(data_window_unit) +@@ -2687,19 +2687,19 @@ class TokenCommNonDeterministic(KernelComponent): + sf_chunk_count = ceil_div(output_sf_bytes, sf_chunk_bytes) + sf_window_unit = ceil_div(sf_chunk_bytes * 2, 3) + reuse_sf_pacing_enabled = ( +- self.token_back_mode == "reuse_dispatch_warps" and sf_window_unit > self.minimum_pacing_window_cycles ++ self.token_back_mode == "reuse_dispatch_warps" and sf_window_unit > self.minimum_pacing_window_cycles # nosec B105 + ) + sf_average_window = Int32(sf_window_unit) + sf_low_window = sf_window_unit * 2 // 5 + sf_high_window = sf_window_unit * 2 + + next_dense_token = global_worker_idx - global_worker_count +- if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): ++ if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 + next_dense_token = Int32(0) + next_dense_token = self.next_token(next_dense_token) + expert_valid_begin = Int32(0) + transfer_phase = Int32(0) +- if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): ++ if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): # nosec B105 + expert_ready_slot_begin = Int32(0) + last_ready_slot = Int32(-1) + +@@ -2707,7 +2707,7 @@ class TokenCommNonDeterministic(KernelComponent): + while local_expert < Int32(self.experts_per_rank): + expert_token_count = owned_sizes[local_expert] + expert_valid_end = expert_valid_begin + expert_token_count +- if cutlass.const_expr(self.token_back_ready_granularity == "expert"): ++ if cutlass.const_expr(self.token_back_ready_granularity == "expert"): # nosec B105 + if next_dense_token < expert_valid_end: + token_tile_count = (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) // Int32( + self.tokens_per_fc1_ready_slot +@@ -2720,7 +2720,7 @@ class TokenCommNonDeterministic(KernelComponent): + + while next_dense_token < expert_valid_end: + token_in_expert = next_dense_token - expert_valid_begin +- if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): ++ if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): # nosec B105 + # Same rank-local expert/tile prefix as FC1 ready; separate counter storage. + ready_slot = expert_ready_slot_begin + token_in_expert // Int32(self.tokens_per_fc1_ready_slot) + if ready_slot != last_ready_slot: +@@ -2765,7 +2765,7 @@ class TokenCommNonDeterministic(KernelComponent): + current_chunk_bytes = Int32(chunk_bytes_this_round) + current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) + stateless_data_pacing_enabled = ( +- self.token_back_mode != "reuse_dispatch_warps" ++ self.token_back_mode != "reuse_dispatch_warps" # nosec B105 + and current_window_unit > self.minimum_pacing_window_cycles + ) + round_start_clock = Int64(0) +@@ -2841,7 +2841,7 @@ class TokenCommNonDeterministic(KernelComponent): + current_chunk_bytes = Int32(chunk_bytes_this_round) + current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) + stateless_sf_pacing_enabled = ( +- self.token_back_mode != "reuse_dispatch_warps" ++ self.token_back_mode != "reuse_dispatch_warps" # nosec B105 + and current_window_unit > self.minimum_pacing_window_cycles + ) + round_start_clock = Int64(0) +@@ -2894,7 +2894,7 @@ class TokenCommNonDeterministic(KernelComponent): + cute.arch.sync_warp() + next_dense_token = self.next_token(next_dense_token) + +- if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): ++ if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): # nosec B105 + # Advance even when this worker skipped the expert (empty experts contribute zero). + expert_ready_slot_begin = expert_ready_slot_begin + ( + expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1) +@@ -2910,9 +2910,9 @@ class TokenCommNonDeterministic(KernelComponent): + # communication CTAs' combine warps, not every persistent CTA's transfer warps. + global_worker_count = self.combine_cta_count * self.combine_warp_count + schedule_counter = None +- if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): ++ if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 + schedule_counter = self._device_workspace.ptr(self.token_back_schedule_region) +- if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): ++ if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 + claimed_token = Int32(0) + if lane_idx == Int32(0): + claimed_token = cute.arch.atomic_add(schedule_counter, Int32(1), sem="relaxed", scope="gpu") +diff --git a/communication/nvlink_domain/token_comm_cta_mainloop.py b/communication/nvlink_domain/token_comm_cta_mainloop.py +index e9cf01a60538d05649f80d79c0cf8cac129566f7..ecc713fcbf9df7cba4ca999b45eee3732167db33 100644 +--- a/communication/nvlink_domain/token_comm_cta_mainloop.py ++++ b/communication/nvlink_domain/token_comm_cta_mainloop.py +@@ -29,7 +29,7 @@ class TokenCommCtaMainloop: + self.dispatch_warp_count = token_comm.dispatch_warp_count + self.combine_warp_start = token_comm.combine_warp_start + self.combine_warp_count = token_comm.combine_warp_count +- if token_comm.token_back_mode != "standalone_warps": ++ if token_comm.token_back_mode != "standalone_warps": # nosec B105 + raise ValueError("Dedicated communication CTAs require standalone token-back warps.") + if not token_back_on_compute_ctas: + dispatch_end = self.dispatch_warp_start + self.dispatch_warp_count +diff --git a/communication/nvlink_domain/token_comm_deterministic.py b/communication/nvlink_domain/token_comm_deterministic.py +index 8b74d44c61c2c1efb3108309c49221e2c2b2588a..af3e8b6545b2c86aafa9ceb2a138cefe41d3d954 100644 +--- a/communication/nvlink_domain/token_comm_deterministic.py ++++ b/communication/nvlink_domain/token_comm_deterministic.py +@@ -137,7 +137,7 @@ class _MetadataPushRouter(KernelComponent): + sorted_metadata_region = "nvlink.token_comm.sorted_metadata" + sorted_scores_region = "nvlink.token_comm.sorted_scores" + pool_expert_base_region = "nvlink.token_comm.pool_expert_base" +- token_src_metadata_region = "nvlink.token_comm.token_src_metadata" ++ token_src_metadata_region = "nvlink.token_comm.token_src_metadata" # nosec B105 + fc1_topk_scores_region = "nvlink.token_comm.fc1_topk_scores" + source_expert_base_region = "nvlink.token_comm.source_expert_base" + push_destination_base_region = "nvlink.token_comm.push_destination_base" +@@ -1347,15 +1347,15 @@ class TokenCommDeterministic(KernelComponent): + fc2_activation_sf_region = "nvlink.token_comm.fc2_activation_sf" + pre_reduced_activation_region = "nvlink.token_comm.pre_reduced_activation" + pre_reduced_activation_sf_region = "nvlink.token_comm.pre_reduced_activation_sf" +- token_back_schedule_region = "nvlink.token_comm.token_back_schedule" ++ token_back_schedule_region = "nvlink.token_comm.token_back_schedule" # nosec B105 + +- token_in_mbarrier_region = "nvlink.token_comm.main_smem.token_in_mbarriers" +- token_back_mbarrier_region = "nvlink.token_comm.main_smem.token_back_mbarriers" ++ token_in_mbarrier_region = "nvlink.token_comm.main_smem.token_in_mbarriers" # nosec B105 ++ token_back_mbarrier_region = "nvlink.token_comm.main_smem.token_back_mbarriers" # nosec B105 + expert_sizes_smem_region = "nvlink.token_comm.main_smem.expert_sizes" +- token_in_activation_smem_region = "nvlink.token_comm.main_smem.token_in_activation" +- token_in_sf_smem_region = "nvlink.token_comm.main_smem.token_in_sf" +- token_back_activation_smem_region = "nvlink.token_comm.main_smem.token_back_activation" +- token_back_sf_smem_region = "nvlink.token_comm.main_smem.token_back_sf" ++ token_in_activation_smem_region = "nvlink.token_comm.main_smem.token_in_activation" # nosec B105 ++ token_in_sf_smem_region = "nvlink.token_comm.main_smem.token_in_sf" # nosec B105 ++ token_back_activation_smem_region = "nvlink.token_comm.main_smem.token_back_activation" # nosec B105 ++ token_back_sf_smem_region = "nvlink.token_comm.main_smem.token_back_sf" # nosec B105 + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: +@@ -1519,7 +1519,7 @@ class TokenCommDeterministic(KernelComponent): + + @property + def token_back_push_data(self) -> bool: +- return self.token_back_mode != "epi_warps" ++ return self.token_back_mode != "epi_warps" # nosec B105 + + @property + def token_back_push_sf(self) -> bool: +@@ -1716,7 +1716,7 @@ class TokenCommDeterministic(KernelComponent): + buffer_space="local", + byte_alignment=128, + ) +- if self.token_back_enabled and self.token_back_schedule_mode == "atomic_counter": ++ if self.token_back_enabled and self.token_back_schedule_mode == "atomic_counter": # nosec B105 + workspace.register( + self.token_back_schedule_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) +@@ -1764,14 +1764,14 @@ class TokenCommDeterministic(KernelComponent): + if not self.token_back_enabled: + return + +- if self.token_back_mode == "standalone_warps": ++ if self.token_back_mode == "standalone_warps": # nosec B105 + token_back_lifetime = workspace.create_overlay( + "nvlink.token_comm.main_smem.standalone_token_back" + ).add_lifetime("token_back") + else: + token_back_lifetime = transfer_overlay.add_lifetime("token_back") + +- if self.token_back_mode == "standalone_warps": ++ if self.token_back_mode == "standalone_warps": # nosec B105 + available_bytes_per_warp = self.standalone_chunk_bytes + else: + activation_bytes = self.bytes_per_token +@@ -1780,13 +1780,13 @@ class TokenCommDeterministic(KernelComponent): + + if self.token_back_push_data: + bytes_per_output_token = self.hidden_size * int(self.combine_format.act_dtype.width) // 8 +- if self.token_back_mode == "standalone_warps": ++ if self.token_back_mode == "standalone_warps": # nosec B105 + chunk_bytes = self.standalone_chunk_bytes + elif available_bytes_per_warp < bytes_per_output_token: + chunk_bytes = self.bytes_per_token + else: + chunk_bytes = bytes_per_output_token +- if self.token_back_mode != "standalone_warps" and bytes_per_output_token % chunk_bytes != 0: ++ if self.token_back_mode != "standalone_warps" and bytes_per_output_token % chunk_bytes != 0: # nosec B105 + raise ValueError("Token-back data chunk bytes must divide one row.") + chunk_elements = chunk_bytes * 8 // int(self.combine_format.act_dtype.width) + token_back_lifetime.register_tensor( +@@ -2098,7 +2098,7 @@ class TokenCommDeterministic(KernelComponent): + ) + token_in_size_barrier.arrive_and_wait() + iket.range_pop() +- if cutlass.const_expr(self.token_back_mode == "standalone_warps"): ++ if cutlass.const_expr(self.token_back_mode == "standalone_warps"): # nosec B105 + sizes_ready_barrier = pipeline.NamedBarrier( + barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count + ) +@@ -2242,7 +2242,7 @@ class TokenCommDeterministic(KernelComponent): + flag_tracker.fire() + cute.arch.sync_warp() + iket.range_pop() +- if cutlass.const_expr(self.token_back_enabled and self.token_back_mode != "standalone_warps"): ++ if cutlass.const_expr(self.token_back_enabled and self.token_back_mode != "standalone_warps"): # nosec B105 + iket.range_push("token_in.transfer_barrier") + transfer_lifetime_barrier = pipeline.NamedBarrier( + barrier_id=self.transfer_lifetime_barrier_id, num_threads=self.transfer_thread_count +@@ -2310,7 +2310,7 @@ class TokenCommDeterministic(KernelComponent): + cute.arch.sync_warp() + + owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) +- if cutlass.const_expr(self.token_back_mode == "standalone_warps"): ++ if cutlass.const_expr(self.token_back_mode == "standalone_warps"): # nosec B105 + sizes_ready_barrier = pipeline.NamedBarrier( + barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count + ) +@@ -2325,7 +2325,7 @@ class TokenCommDeterministic(KernelComponent): + activation_chunk_count = ceil_div(output_token_bytes, activation_chunk_bytes) + data_window_unit = ceil_div(activation_chunk_bytes * 2, 3) + reuse_data_pacing_enabled = ( +- self.token_back_mode == "reuse_dispatch_warps" and data_window_unit > self.minimum_pacing_window_cycles ++ self.token_back_mode == "reuse_dispatch_warps" and data_window_unit > self.minimum_pacing_window_cycles # nosec B105 + ) + # Preserve the empirical low:initial:high ratio of 1:2.5:5. + data_average_window = Int32(data_window_unit) +@@ -2338,14 +2338,14 @@ class TokenCommDeterministic(KernelComponent): + sf_chunk_count = ceil_div(output_sf_bytes, sf_chunk_bytes) + sf_window_unit = ceil_div(sf_chunk_bytes * 2, 3) + reuse_sf_pacing_enabled = ( +- self.token_back_mode == "reuse_dispatch_warps" and sf_window_unit > self.minimum_pacing_window_cycles ++ self.token_back_mode == "reuse_dispatch_warps" and sf_window_unit > self.minimum_pacing_window_cycles # nosec B105 + ) + sf_average_window = Int32(sf_window_unit) + sf_low_window = sf_window_unit * 2 // 5 + sf_high_window = sf_window_unit * 2 + + next_dense_token = global_worker_idx - global_worker_count +- if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): ++ if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 + next_dense_token = Int32(0) + next_dense_token = self.next_token(next_dense_token) + expert_valid_begin = Int32(0) +@@ -2401,7 +2401,7 @@ class TokenCommDeterministic(KernelComponent): + current_chunk_bytes = Int32(chunk_bytes_this_round) + current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) + stateless_data_pacing_enabled = ( +- self.token_back_mode != "reuse_dispatch_warps" ++ self.token_back_mode != "reuse_dispatch_warps" # nosec B105 + and current_window_unit > self.minimum_pacing_window_cycles + ) + round_start_clock = Int64(0) +@@ -2468,7 +2468,7 @@ class TokenCommDeterministic(KernelComponent): + current_chunk_bytes = Int32(chunk_bytes_this_round) + current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) + stateless_sf_pacing_enabled = ( +- self.token_back_mode != "reuse_dispatch_warps" ++ self.token_back_mode != "reuse_dispatch_warps" # nosec B105 + and current_window_unit > self.minimum_pacing_window_cycles + ) + round_start_clock = Int64(0) +@@ -2520,9 +2520,9 @@ class TokenCommDeterministic(KernelComponent): + def next_token(self, current_token: Int32) -> Int32: + global_worker_count = self.promised_launchable_sm_count * self.transfer_warp_count + schedule_counter = None +- if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): ++ if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 + schedule_counter = self._device_workspace.ptr(self.token_back_schedule_region) +- if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): ++ if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 + claimed_token = Int32(0) + if self._lane_idx == Int32(0): + claimed_token = cute.arch.atomic_add(schedule_counter, Int32(1), sem="relaxed", scope="gpu") +diff --git a/helpers/topk_reduce.py b/helpers/topk_reduce.py +index 6e4d0fe7d363b8554e5b2789240e555ccc2499ea..ce4cb6465a0fb430bcb2a93acf2d8a158ff2989b 100644 +--- a/helpers/topk_reduce.py ++++ b/helpers/topk_reduce.py +@@ -38,7 +38,7 @@ from .dsl_helpers import mark_alignment + # fp4 (e2m1) -> fp32 register decode. + # + # Blackwell has no e2m1->f32 upconvert: the framework's ``term.load().to(f32)`` +-# lowers to an ALU subnormal-normalization path (~60% DRAM SOL). Both helpers ++# lowers to an ALU subnormal-normalization path. Both helpers + # below force a table-driven decode instead; ``e2m1_reg`` (N e2m1 codes, N % 8 + # == 0) is read as packed b32 words and N fp32 values are written into + # ``fp32_reg`` in code order. The 16 e2m1 values are exact in fp32, so the two +diff --git a/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py b/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py +index cedc19ea477a24fd5d58e14777752fe84e44bd85..8057166cabff92d158798311a3fc7b1549d408b7 100644 +--- a/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py ++++ b/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py +@@ -959,6 +959,8 @@ class SwapABGatedActEpilogue(KernelComponent): + "intermediate_gateup_size": StaticOrRuntimeIntegerType, + "combine_format": CombineFormat, + "gate_up_clamp": Optional[float], ++ "swiglu_alpha": OptionalRequirement(Optional[float]), ++ "swiglu_beta": OptionalRequirement(Optional[float]), + # SiTU (Kimi K3) selects a different gated-activation core; optional so + # existing SwiGLU descriptors stay valid unchanged. See _resolve_situ_betas. + "situ_beta": OptionalRequirement(Optional[float]), +@@ -999,6 +1001,8 @@ class SwapABGatedActEpilogue(KernelComponent): + raise ValueError(f"SiTU beta parameters must be positive, got {beta} and {linear_beta}.") + if problem_desc["gate_up_clamp"] is not None: + raise ValueError("SiTU does not support gate_up_clamp; the two activation variants are exclusive.") ++ if problem_desc.get("swiglu_alpha") is not None: ++ raise ValueError("SiTU does not support SwiGLUBias constants.") + return float(beta), float(linear_beta) + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: +@@ -1010,6 +1014,13 @@ class SwapABGatedActEpilogue(KernelComponent): + self.intermediate_gateup_size = problem_desc["intermediate_gateup_size"] + self.combine_format = problem_desc["combine_format"] + self.gate_up_clamp = problem_desc["gate_up_clamp"] ++ self.swiglu_alpha = problem_desc.get("swiglu_alpha") ++ self.swiglu_beta = problem_desc.get("swiglu_beta") ++ if (self.swiglu_alpha is None) != (self.swiglu_beta is None): ++ raise ValueError("swiglu_alpha and swiglu_beta must be set together.") ++ if self.swiglu_alpha is not None: ++ self.swiglu_alpha = float(self.swiglu_alpha) ++ self.swiglu_beta = float(self.swiglu_beta) + self.situ_beta, self.situ_linear_beta = self._resolve_situ_betas(problem_desc) + self.mma_tiler_mnk = impl_desc["mma_tiler_mnk"] + self.cluster_shape_mn = impl_desc["cluster_shape_mn"] +@@ -1765,7 +1776,8 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): + # gate = min(gate, +limit) (upper bound only) + # up = clamp(up, -limit, +limit) (symmetric) + # 3. gated activation, either SwiGLU or SiTU (mutually exclusive; SiTU has no clamp): +- # SwiGLU: out = up * gate * sigmoid(gate) ++ # SwiGLU: out = (up + swiglu_beta) * gate * sigmoid(swiglu_alpha * gate) ++ # Standard SwiGLU uses swiglu_alpha=1 and swiglu_beta=0. + # SiTU: out = beta*tanh(gate/beta)*sigmoid(gate) * linear_beta*tanh(up/linear_beta) + # sigmoid(x) = rcp(1 + exp2(-x * log2e)) + # +@@ -1777,8 +1789,12 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): + out = cute.make_rmem_tensor((n,), cutlass.Float32) + log2_e = 1.4426950408889634 + +- neg_log2e_pair = (cutlass.Float32(-log2_e), cutlass.Float32(-log2_e)) ++ activation_alpha = 1.0 if self.swiglu_alpha is None else self.swiglu_alpha ++ neg_activation_alpha_log2e = -activation_alpha * log2_e ++ neg_log2e_pair = (cutlass.Float32(neg_activation_alpha_log2e), cutlass.Float32(neg_activation_alpha_log2e)) + one_pair = (cutlass.Float32(1.0), cutlass.Float32(1.0)) ++ if cutlass.const_expr(self.swiglu_beta is not None): ++ beta_pair = (cutlass.Float32(self.swiglu_beta), cutlass.Float32(self.swiglu_beta)) + if cutlass.const_expr(self.gate_up_clamp is not None): + limit = cutlass.Float32(self.gate_up_clamp) + +@@ -1854,7 +1870,10 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): + sigmoid_pair = _sigmoid(g0, g1) + + if cutlass.const_expr(self.situ_beta is None): +- # SwiGLU: out = up * gate * sigmoid(gate) ++ # SwiGLU: out = (up + swiglu_beta) * gate * sigmoid(swiglu_alpha * gate) ++ # The SwiGLUBias offset follows the up clamp. ++ if cutlass.const_expr(self.swiglu_beta is not None): ++ u0, u1 = cute.arch.add_packed_f32x2((u0, u1), beta_pair) + ug = cute.arch.mul_packed_f32x2((u0, u1), (g0, g1)) + out_pair = cute.arch.mul_packed_f32x2(ug, sigmoid_pair) + else: +diff --git a/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py b/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py +index d7f4e53daff8e7e2cee8444b21c57f114c9a8fe1..7dad5c86e02b22f49021cb39c600d41f65243cd1 100644 +--- a/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py ++++ b/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py +@@ -73,6 +73,8 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): + "b_major_mode": OperandMajorMode, + "combine_format": CombineFormat, + "gate_up_clamp": Optional[float], ++ "swiglu_alpha": OptionalRequirement(Optional[float]), ++ "swiglu_beta": OptionalRequirement(Optional[float]), + "situ_beta": OptionalRequirement(Optional[float]), + "situ_linear_beta": OptionalRequirement(Optional[float]), + "world_size": int, +@@ -127,6 +129,8 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): + self.b_major_mode = problem_desc["b_major_mode"] + self.combine_format = problem_desc["combine_format"] + self.gate_up_clamp = problem_desc["gate_up_clamp"] ++ self.swiglu_alpha = problem_desc.get("swiglu_alpha") ++ self.swiglu_beta = problem_desc.get("swiglu_beta") + self.situ_beta = problem_desc.get("situ_beta") + self.situ_linear_beta = problem_desc.get("situ_linear_beta") + self.world_size = problem_desc["world_size"] +@@ -168,8 +172,8 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): + # static M is required only when an S>0 helper boundary must be + # compiled into the scheduler gate. + self.home_expert_count = self.local_expert_count +- self.threads_per_cta = 16 * 32 if self.token_back_mode == "standalone_warps" else 12 * 32 +- self.other_warp_register_count = 64 if self.token_back_mode == "standalone_warps" else 72 ++ self.threads_per_cta = 16 * 32 if self.token_back_mode == "standalone_warps" else 12 * 32 # nosec B105 ++ self.other_warp_register_count = 64 if self.token_back_mode == "standalone_warps" else 72 # nosec B105 + + self.mixed_cga_config = NonClcMixedCgaConfig( + preferred_cluster_shape=self.cluster_shape_mn, +@@ -453,6 +457,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): + f"epiflag{epi_flags}_clusters{self.launch_cluster_count}_ctas{self.mixed_cga_config.total_cta_cnt}_" + f"tokeninflag{self.token_in_flag_batch}_tokenback{self.token_back_mode}_" + f"combine{self.combine_format.name}_clamp{self.gate_up_clamp}_" ++ f"swiglua{self.swiglu_alpha}_swiglub{self.swiglu_beta}_" + # The SiTU betas are codegen-time constants, so they must key the compiled kernel. + f"situ{self.situ_beta}x{self.situ_linear_beta}_" + f"{'apply_topk_fc1' if self.apply_topk_at_fc1 else 'apply_topk_fc2'}_" +@@ -1032,7 +1037,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): + pipeline_init_wait_mixed_cga( + self.cluster_shape_mn, self.resolved_fallback_cluster_shape_mn, is_fallback_cluster + ) +- finalize_barrier_thread_count = 12 * 32 if self.token_back_mode == "standalone_warps" else 8 * 32 ++ finalize_barrier_thread_count = 12 * 32 if self.token_back_mode == "standalone_warps" else 8 * 32 # nosec B105 + finalize_barrier = pipeline.NamedBarrier(barrier_id=13, num_threads=finalize_barrier_thread_count) + + if warp_idx == self.scheduler_warp_id: +@@ -1161,7 +1166,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): + self.token_comm.token_in(self._smem_workspace, smem_base) + iket.range_pop() + if cutlass.const_expr( +- self.token_comm.token_back_enabled and self.token_back_mode != "standalone_warps" ++ self.token_comm.token_back_enabled and self.token_back_mode != "standalone_warps" # nosec B105 + ): + iket.range_push("mega.token_back") + self.token_comm.token_back(self._smem_workspace, smem_base) +@@ -1170,7 +1175,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): + iket.range_push("mega.tail_reset") + self.token_comm.reset_tail() + iket.range_pop() +- elif cutlass.const_expr(self.token_back_mode == "standalone_warps"): ++ elif cutlass.const_expr(self.token_back_mode == "standalone_warps"): # nosec B105 + iket.range_push("mega.token_back_standalone") + self.token_comm.token_back(self._smem_workspace, smem_base) + iket.range_pop() +diff --git a/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py b/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py +index 03507e3bb60f159fca2343204339810a3db2d14d..7fbb2be22b7645c0fe116b4db662c64a053bda65 100644 +--- a/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py ++++ b/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py +@@ -63,6 +63,13 @@ class SwapABGatedActEpilogue(_SharedSwapABGatedActEpilogue): + self.intermediate_gateup_size = problem_desc["intermediate_gateup_size"] + self.combine_format = problem_desc["combine_format"] + self.gate_up_clamp = problem_desc["gate_up_clamp"] ++ self.swiglu_alpha = problem_desc.get("swiglu_alpha") ++ self.swiglu_beta = problem_desc.get("swiglu_beta") ++ if (self.swiglu_alpha is None) != (self.swiglu_beta is None): ++ raise ValueError("swiglu_alpha and swiglu_beta must be set together.") ++ if self.swiglu_alpha is not None: ++ self.swiglu_alpha = float(self.swiglu_alpha) ++ self.swiglu_beta = float(self.swiglu_beta) + self.situ_beta, self.situ_linear_beta = self._resolve_situ_betas(problem_desc) + self.mma_tiler_mnk = impl_desc["mma_tiler_mnk"] + self.cluster_shape_mn = impl_desc["cluster_shape_mn"] +diff --git a/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_local_mega_moe_kernel.py b/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_local_mega_moe_kernel.py +index a486110e0e3eb4b673d5857f7d08484cc24893a5..3e86b1ad34e2e31ab8b3a394208e921156c4b552 100644 +--- a/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_local_mega_moe_kernel.py ++++ b/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_local_mega_moe_kernel.py +@@ -67,6 +67,8 @@ class BlockScaledSwapAbLocalMegaMoeKernel(KernelClass): + "a_major_mode": OperandMajorMode, + "b_major_mode": OperandMajorMode, + "gate_up_clamp": Optional[float], ++ "swiglu_alpha": OptionalRequirement(Optional[float]), ++ "swiglu_beta": OptionalRequirement(Optional[float]), + "situ_beta": OptionalRequirement(Optional[float]), + "situ_linear_beta": OptionalRequirement(Optional[float]), + "topk": int, +@@ -116,6 +118,8 @@ class BlockScaledSwapAbLocalMegaMoeKernel(KernelClass): + self.b_major_mode = problem_desc["b_major_mode"] + self.combine_format = CombineFormat.parse("bf16") + self.gate_up_clamp = problem_desc["gate_up_clamp"] ++ self.swiglu_alpha = problem_desc.get("swiglu_alpha") ++ self.swiglu_beta = problem_desc.get("swiglu_beta") + self.situ_beta = problem_desc.get("situ_beta") + self.situ_linear_beta = problem_desc.get("situ_linear_beta") + self.topk = problem_desc["topk"] +@@ -410,6 +414,7 @@ class BlockScaledSwapAbLocalMegaMoeKernel(KernelClass): + f"fc1epiflag{self.fc1_epi_flag_batch}_clusters{self.launch_cluster_count}_" + f"routingflag{self.routing_flag_batch}_" + f"combine{self.combine_format.name}_clamp{self.gate_up_clamp}_" ++ f"swiglua{self.swiglu_alpha}_swiglub{self.swiglu_beta}_" + # The SiTU betas are codegen-time constants, so they must key the compiled kernel. + f"situ{self.situ_beta}x{self.situ_linear_beta}_" + f"{'apply_topk_fc1' if self.apply_topk_at_fc1 else 'apply_topk_fc2'}_" +diff --git a/kernel_src/rubin/inference/local_mega/local_routing.py b/kernel_src/rubin/inference/local_mega/local_routing.py +index ce7877ec231429f6487668527285388c09bea0bd..fd3ceb087e89c02c041b0e794dc22fbc15a1e67b 100644 +--- a/kernel_src/rubin/inference/local_mega/local_routing.py ++++ b/kernel_src/rubin/inference/local_mega/local_routing.py +@@ -174,7 +174,7 @@ class _MetadataPushRouter(KernelComponent): + sorted_metadata_region = "local_mega.routing.sorted_metadata" + sorted_scores_region = "local_mega.routing.sorted_scores" + pool_expert_base_region = "local_mega.routing.pool_expert_base" +- token_src_metadata_region = "local_mega.routing.token_src_metadata" ++ token_src_metadata_region = "local_mega.routing.token_src_metadata" # nosec B105 + fc1_topk_scores_region = "local_mega.routing.fc1_topk_scores" + source_expert_base_region = "local_mega.routing.source_expert_base" + sorted_metadata_ready_region = "local_mega.routing.sorted_metadata_ready" +@@ -1021,9 +1021,9 @@ class LocalRouting(KernelComponent): + fc1_activation_sf_region = "local_mega.routing.fc1_activation_sf" + pre_reduced_activation_region = "local_mega.routing.pre_reduced_activation" + +- token_in_mbarrier_region = "local_mega.routing.main_smem.token_in_mbarriers" ++ token_in_mbarrier_region = "local_mega.routing.main_smem.token_in_mbarriers" # nosec B105 + expert_sizes_smem_region = "local_mega.routing.main_smem.expert_sizes" +- token_in_sf_smem_region = "local_mega.routing.main_smem.token_in_sf" ++ token_in_sf_smem_region = "local_mega.routing.main_smem.token_in_sf" # nosec B105 + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: +diff --git a/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py b/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py +index bd269d76d534d195778186d3844982101ee87860..f0868f743ba480f8d3c8f0a9aac7bf31cdc90daa 100644 +--- a/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py ++++ b/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py +@@ -896,6 +896,8 @@ class SwapABGatedActEpilogue(KernelComponent): + "intermediate_gateup_size": StaticOrRuntimeIntegerType, + "combine_format": CombineFormat, + "gate_up_clamp": Optional[float], ++ "swiglu_alpha": OptionalRequirement(Optional[float]), ++ "swiglu_beta": OptionalRequirement(Optional[float]), + # SiTU (Kimi K3) selects a different gated-activation core; optional so + # existing SwiGLU descriptors stay valid unchanged. See _resolve_situ_betas. + "situ_beta": OptionalRequirement(Optional[float]), +@@ -937,6 +939,8 @@ class SwapABGatedActEpilogue(KernelComponent): + raise ValueError(f"SiTU beta parameters must be positive, got {beta} and {linear_beta}.") + if problem_desc["gate_up_clamp"] is not None: + raise ValueError("SiTU does not support gate_up_clamp; the two activation variants are exclusive.") ++ if problem_desc.get("swiglu_alpha") is not None: ++ raise ValueError("SiTU does not support SwiGLUBias constants.") + return float(beta), float(linear_beta) + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: +@@ -948,6 +952,13 @@ class SwapABGatedActEpilogue(KernelComponent): + self.intermediate_gateup_size = problem_desc["intermediate_gateup_size"] + self.combine_format = problem_desc["combine_format"] + self.gate_up_clamp = problem_desc["gate_up_clamp"] ++ self.swiglu_alpha = problem_desc.get("swiglu_alpha") ++ self.swiglu_beta = problem_desc.get("swiglu_beta") ++ if (self.swiglu_alpha is None) != (self.swiglu_beta is None): ++ raise ValueError("swiglu_alpha and swiglu_beta must be set together.") ++ if self.swiglu_alpha is not None: ++ self.swiglu_alpha = float(self.swiglu_alpha) ++ self.swiglu_beta = float(self.swiglu_beta) + self.situ_beta, self.situ_linear_beta = self._resolve_situ_betas(problem_desc) + self.mma_tiler_mnk = impl_desc["mma_tiler_mnk"] + self.cluster_shape_mn = impl_desc["cluster_shape_mn"] +@@ -1639,7 +1650,8 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): + # gate = min(gate, +limit) (upper bound only) + # up = clamp(up, -limit, +limit) (symmetric) + # 3. gated activation, either SwiGLU or SiTU (mutually exclusive; SiTU has no clamp): +- # SwiGLU: out = up * gate * sigmoid(gate) ++ # SwiGLU: out = (up + swiglu_beta) * gate * sigmoid(swiglu_alpha * gate) ++ # Standard SwiGLU uses swiglu_alpha=1 and swiglu_beta=0. + # SiTU: out = beta*tanh(gate/beta)*sigmoid(gate) * linear_beta*tanh(up/linear_beta) + # sigmoid(x) = rcp(1 + exp2(-x * log2e)) + # +@@ -1651,8 +1663,12 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): + out = cute.make_rmem_tensor((n,), cutlass.Float32) + log2_e = 1.4426950408889634 + +- neg_log2e_pair = (cutlass.Float32(-log2_e), cutlass.Float32(-log2_e)) ++ activation_alpha = 1.0 if self.swiglu_alpha is None else self.swiglu_alpha ++ neg_activation_alpha_log2e = -activation_alpha * log2_e ++ neg_log2e_pair = (cutlass.Float32(neg_activation_alpha_log2e), cutlass.Float32(neg_activation_alpha_log2e)) + one_pair = (cutlass.Float32(1.0), cutlass.Float32(1.0)) ++ if cutlass.const_expr(self.swiglu_beta is not None): ++ beta_pair = (cutlass.Float32(self.swiglu_beta), cutlass.Float32(self.swiglu_beta)) + if cutlass.const_expr(self.gate_up_clamp is not None): + limit = cutlass.Float32(self.gate_up_clamp) + +@@ -1728,7 +1744,10 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): + sigmoid_pair = _sigmoid(g0, g1) + + if cutlass.const_expr(self.situ_beta is None): +- # SwiGLU: out = up * gate * sigmoid(gate) ++ # SwiGLU: out = (up + swiglu_beta) * gate * sigmoid(swiglu_alpha * gate) ++ # The SwiGLUBias offset follows the up clamp. ++ if cutlass.const_expr(self.swiglu_beta is not None): ++ u0, u1 = cute.arch.add_packed_f32x2((u0, u1), beta_pair) + ug = cute.arch.mul_packed_f32x2((u0, u1), (g0, g1)) + out_pair = cute.arch.mul_packed_f32x2(ug, sigmoid_pair) + else: +@@ -1985,7 +2004,7 @@ class SwapABFc2Epilogue(_ImmutableAfterInit): + publish: cutlass.Constexpr = self.token_back_enabled + flag_address = Int64(0) + if cutlass.const_expr(publish): +- if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): ++ if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): # nosec B105 + # The scheduler maps every fallback CTA into the preferred + # logical cluster. N=1 keeps this slot common to all hidden + # shards; padded hidden CTAs must also publish their completion. +diff --git a/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py b/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py +index f49de382daeef5a9751a19dd190008f02827c54c..3c5bb77f999cd0013dd68ec11668adefe816727d 100644 +--- a/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py ++++ b/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py +@@ -75,6 +75,8 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): + "b_major_mode": OperandMajorMode, + "combine_format": CombineFormat, + "gate_up_clamp": Optional[float], ++ "swiglu_alpha": OptionalRequirement(Optional[float]), ++ "swiglu_beta": OptionalRequirement(Optional[float]), + "situ_beta": OptionalRequirement(Optional[float]), + "situ_linear_beta": OptionalRequirement(Optional[float]), + "world_size": int, +@@ -136,6 +138,8 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): + self.b_major_mode = problem_desc["b_major_mode"] + self.combine_format = problem_desc["combine_format"] + self.gate_up_clamp = problem_desc["gate_up_clamp"] ++ self.swiglu_alpha = problem_desc.get("swiglu_alpha") ++ self.swiglu_beta = problem_desc.get("swiglu_beta") + self.situ_beta = problem_desc.get("situ_beta") + self.situ_linear_beta = problem_desc.get("situ_linear_beta") + self.world_size = problem_desc["world_size"] +@@ -202,8 +206,8 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): + # static M is required only when an S>0 helper boundary must be + # compiled into the scheduler gate. + self.home_expert_count = self.local_expert_count +- self.threads_per_cta = 16 * 32 if self.token_back_mode == "standalone_warps" else 12 * 32 +- self.other_warp_register_count = 64 if self.token_back_mode == "standalone_warps" else 72 ++ self.threads_per_cta = 16 * 32 if self.token_back_mode == "standalone_warps" else 12 * 32 # nosec B105 ++ self.other_warp_register_count = 64 if self.token_back_mode == "standalone_warps" else 72 # nosec B105 + + self.mixed_cga_config = NonClcMixedCgaConfig( + preferred_cluster_shape=self.cluster_shape_mn, +@@ -239,7 +243,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): + sf_vec_size=self.sf_vec_size, + ) + tmem_plan = make_tcgen05_tmem_plan(mma_instruction, self.architecture, self.mma_tiler_mnk) +- if self.token_back_ready_granularity == "token_tile": ++ if self.token_back_ready_granularity == "token_tile": # nosec B105 + tokens_per_fc1_ready_slot, fc2_done_signals_per_token_tile = derive_fc2_ready_geometry( + hidden_size=self.hidden_size, + mma_tiler_mnk=self.mma_tiler_mnk, +@@ -301,7 +305,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): + ) + self.scheduler = scheduler_type(fc12_problem_desc, resolved_impl_desc) + self.epilogue = SwapABGatedActEpilogue(fc12_problem_desc, resolved_impl_desc) +- if self.token_back_ready_granularity == "token_tile": ++ if self.token_back_ready_granularity == "token_tile": # nosec B105 + self._validate_token_tile_ready_contract() + if self.epilogue.fc1_output_dtype is not self.b_dtype: + raise ValueError("Epilogue FC1 output dtype must match the Mainloop B dtype.") +@@ -366,7 +370,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): + or any(isinstance(w, bool) or not isinstance(w, int) or w <= 0 for w in warps) + ): + raise ValueError(f"A split needs comm_warps as two positive ints, got {warps}.") +- if self.token_back_mode != "standalone_warps": ++ if self.token_back_mode != "standalone_warps": # nosec B105 + raise ValueError("A CTA split requires token_back_mode='standalone_warps'.") + dispatch_warps, back_warps = warps + # Check the warps that actually coexist in each concrete CTA branch. +@@ -438,7 +442,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): + def _validate_geometry(self) -> None: + if self.token_back_ready_granularity not in ("expert", "token_tile"): + raise ValueError(f"Unsupported token_back_ready_granularity {self.token_back_ready_granularity!r}.") +- if self.token_back_ready_granularity == "token_tile": ++ if self.token_back_ready_granularity == "token_tile": # nosec B105 + if not ( + self.schedule_mode in ("grouped", "phase_interleave") + and self.work_id_mode == "atomic_counter" +@@ -678,12 +682,13 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): + f"epiflag{epi_flags}_clusters{self.launch_cluster_count}_ctas{self.mixed_cga_config.total_cta_cnt}_" + f"tokeninflag{self.token_in_flag_batch}_tokenback{self.token_back_mode}_" + f"combine{self.combine_format.name}_clamp{self.gate_up_clamp}_" ++ f"swiglua{self.swiglu_alpha}_swiglub{self.swiglu_beta}_" + # The SiTU betas are codegen-time constants, so they must key the compiled kernel. + f"situ{self.situ_beta}x{self.situ_linear_beta}_" + f"{'apply_topk_fc1' if self.apply_topk_at_fc1 else 'apply_topk_fc2'}_" + f"{'inkernel_redg' if self.reduce_topk_in_kernel else 'separate_reduce'}" + "_dlb_v1" +- + ("_fc2readytile_v1" if self.token_back_ready_granularity == "token_tile" else "") ++ + ("_fc2readytile_v1" if self.token_back_ready_granularity == "token_tile" else "") # nosec B105 + ) + + def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): +@@ -1275,7 +1280,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): + pipeline_init_wait_mixed_cga( + self.cluster_shape_mn, self.resolved_fallback_cluster_shape_mn, is_fallback_cluster + ) +- finalize_barrier_thread_count = 12 * 32 if self.token_back_mode == "standalone_warps" else 8 * 32 ++ finalize_barrier_thread_count = 12 * 32 if self.token_back_mode == "standalone_warps" else 8 * 32 # nosec B105 + finalize_barrier = pipeline.NamedBarrier(barrier_id=13, num_threads=finalize_barrier_thread_count) + + if warp_idx == self.scheduler_warp_id: +@@ -1405,7 +1410,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): + self.token_comm.token_in(self._smem_workspace, smem_base) + iket.range_pop() + if cutlass.const_expr( +- self.token_comm.token_back_enabled and self.token_back_mode != "standalone_warps" ++ self.token_comm.token_back_enabled and self.token_back_mode != "standalone_warps" # nosec B105 + ): + iket.range_push("mega.token_back") + self.token_comm.token_back(self._smem_workspace, smem_base) +@@ -1414,7 +1419,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): + iket.range_push("mega.tail_reset") + self.token_comm.reset_tail() + iket.range_pop() +- elif cutlass.const_expr(self.token_back_mode == "standalone_warps"): ++ elif cutlass.const_expr(self.token_back_mode == "standalone_warps"): # nosec B105 + iket.range_push("mega.token_back_standalone") + self.token_comm.token_back(self._smem_workspace, smem_base) + iket.range_pop() +diff --git a/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py b/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py +index 96b8342673f6943862b81bd9ea8e3aa85d1fac30..c5567f3df5a0031dd77a9c038d9b5b03f510ebd8 100644 +--- a/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py ++++ b/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py +@@ -109,6 +109,8 @@ class BlockScaledSwapAbGenphaseMoeKernel(KernelClass): + "quant_kind": str, + "combine_format": CombineFormat, + "gate_up_clamp": Optional[float], ++ "swiglu_alpha": OptionalRequirement(Optional[float]), ++ "swiglu_beta": OptionalRequirement(Optional[float]), + "world_size": int, + "topk": int, + "topk_index_dtype": type, +@@ -169,6 +171,8 @@ class BlockScaledSwapAbGenphaseMoeKernel(KernelClass): + self.acc_dtype = tcgen05_block_scaled_acc_dtype + self.combine_format = problem_desc["combine_format"] + self.gate_up_clamp = problem_desc["gate_up_clamp"] ++ self.swiglu_alpha = problem_desc.get("swiglu_alpha") ++ self.swiglu_beta = problem_desc.get("swiglu_beta") + self.world_size = problem_desc["world_size"] + self.topk = problem_desc["topk"] + self.topk_index_dtype = problem_desc["topk_index_dtype"] +@@ -546,6 +550,7 @@ class BlockScaledSwapAbGenphaseMoeKernel(KernelClass): + f"epiflag{epi_flags}_clusters{self.launch_cluster_count}_" + f"pusher{self.pusher_cta_count}_refine{self.refine_participant_groups}x{self.refine_output_stages}_" + f"combine{self.combine_format.name}_clamp{self.gate_up_clamp}_" ++ f"swiglua{self.swiglu_alpha}_swiglub{self.swiglu_beta}_" + f"{'apply_topk_fc1' if self.apply_topk_at_fc1 else 'apply_topk_fc2'}_" + f"{'inkernel_reduce' if self.reduce_topk_in_kernel else 'separate_reduce'}_" + f"abstages{self._mainloop.num_weight_ab_stages}x{self._mainloop.num_token_ab_stages}" +diff --git a/kernel_src/schedulers/fc12_scheduler.py b/kernel_src/schedulers/fc12_scheduler.py +index b014b25f72a48e6be38922236faa0a84c9672a1a..9fd349cd80481271f08825c66ae28c53313ba93f 100644 +--- a/kernel_src/schedulers/fc12_scheduler.py ++++ b/kernel_src/schedulers/fc12_scheduler.py +@@ -906,8 +906,6 @@ class PhaseInterleavedFc12Scheduler(SchedulerBase): + minimum_claim_count, task_mapping_state + ) + spin_wait(fc1_counter_pointer, lambda value: value >= claim_target) +- if stream_ends_before_target: +- fc1_exhausted = Boolean(True) + return fc1_exhausted + + @cute.jit +diff --git a/kernel_src/schedulers/non_clc_mixed_cga.py b/kernel_src/schedulers/non_clc_mixed_cga.py +index 841157953fa0eac1ab379feecf0aed589055bc14..ba879cc6ac4d11bacaa4341f6e2f08765a4b1633 100644 +--- a/kernel_src/schedulers/non_clc_mixed_cga.py ++++ b/kernel_src/schedulers/non_clc_mixed_cga.py +@@ -134,7 +134,7 @@ class NonClcMixedCgaSchedulerWorker: + cluster_broadcast_region = "non_clc_mixed_cga.cluster_broadcast" + work_id_counter_region = "non_clc_mixed_cga.work_id_counters" + fallback_registration_counter_region = "non_clc_mixed_cga.fallback_registration_counter" +- fallback_group_token_region = "non_clc_mixed_cga.fallback_group_tokens" ++ fallback_group_token_region = "non_clc_mixed_cga.fallback_group_tokens" # nosec B105 + + def __init__( + self, diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 18640486388d..23dfb7a6c9ac 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -1,5 +1,26 @@ schema_version: 1 vendors: + cutedsl-megamoe: + url: https://gitlab-master.nvidia.com/bangyus/cutedsl_megamoe.git + branch: main + commit: 382ce4a6f111d0769f31c4c975e1c643f0dc4f70 + source: sources + destination: tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe + include: + - '*.py' + - communication/**/*.py + - helpers/[!mt]*.py + - helpers/megamoe_aot.py + - helpers/topk_reduce.py + - kernel_src/*.py + - kernel_src/blackwell/*.py + - kernel_src/rubin/*.py + - kernel_src/*/inference/**/*.py + - kernel_src/schedulers/[!w]*.py + - kernel_src/schedulers/work_id_claim.py + patch: 3rdparty/vendor_patches/cutedsl-megamoe.patch + patch_digest: sha256:0676cec5fe7c087f93bb6aec9d15324460660e74f3caa94a86d7fcb7253deda8 + digest: sha256-tree-v1:bda2e41e83acec2801ecb5020b4fadbab50efb051c1911b5a204e5c0efdba96d flashinfer-kda-bt16: url: https://github.com/flashinfer-ai/flashinfer.git branch: main @@ -23,6 +44,18 @@ vendors: patch: 3rdparty/vendor_patches/flashinfer-prims-ts.patch patch_digest: sha256:fc5c23ec3d65c562025a1845f211fdac66bfcd3da156dbb8385f4c4407d20d07 digest: sha256-tree-v1:f2f444f7e6d4ad688fbf4e1e9eed11ff7c78d835cfd0af5ef9d7660780b059e0 + megamoe-scheduler: + url: https://gitlab-master.nvidia.com/jintaop/cutedsl_eplb_scheduler_copy.git + branch: trtllm/pr-19646-single-plane-spec + commit: 4dc8616031b18dd948bdf11b49afe0d86be15e22 + source: megamoe_scheduler + destination: tensorrt_llm/_torch/moe/kernels/megamoe_scheduler + include: + - __init__.py + - cuda_scheduler/*.py + - geometry.py + - sami/*.py + digest: sha256-tree-v1:00210684fbffaa623fc00839bcacce000650db9090cf7a1a07bdd0ef3a9afc33 openengine-proto: url: https://github.com/ai-dynamo/openengine.git tag: v0.1.0 diff --git a/cpp/tensorrt_llm/kernels/moe/CMakeLists.txt b/cpp/tensorrt_llm/kernels/moe/CMakeLists.txt index 2c4a27545e16..e5bc72b46f7d 100644 --- a/cpp/tensorrt_llm/kernels/moe/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/moe/CMakeLists.txt @@ -20,6 +20,7 @@ # into kernels_src, exactly as it did before the move. add_subdirectory(cuteDsl) # cute_dsl_src add_subdirectory(trtllmGen) # trtllm_gen_fp8_block_scale_moe{,_routing} +add_subdirectory(loadBalance/dynamicEplb) # moe_dynamic_eplb_src # cutlass/: moe_gemm_src and its per-arch OBJECT libraries are defined in # ../cutlass_kernels/CMakeLists.txt, next to the CUTLASS instantiation generator # they depend on. @@ -28,7 +29,9 @@ add_subdirectory(trtllmGen) # trtllm_gen_fp8_block_scale_moe{,_routing} # kernels_src, so the backend directories are named here, next to the # add_subdirectory() calls they correspond to, and nowhere else. file(GLOB_RECURSE MOE_SRC *.cpp *.cu) -list(FILTER MOE_SRC EXCLUDE REGEX "/kernels/moe/(cutlass|trtllmGen|cuteDsl)/.*") +list(FILTER MOE_SRC EXCLUDE REGEX + "/kernels/moe/(cutlass|trtllmGen|cuteDsl|loadBalance/dynamicEplb)/.*") +list(APPEND MOE_SRC $) set(MOE_KERNELS_SRC ${MOE_SRC} PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/CMakeLists.txt b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/CMakeLists.txt new file mode 100644 index 000000000000..1edbb9c4c973 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/CMakeLists.txt @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not +# use this file except in compliance with the License. You may obtain a copy of +# the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations under +# the License. + +set(DYNAMIC_EPLB_ARCHS) +foreach(ARCH 100 103 107) + if("${ARCH}" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG) + list(APPEND DYNAMIC_EPLB_ARCHS ${ARCH}) + endif() +endforeach() + +set(DYNAMIC_EPLB_SOURCES) +set(DYNAMIC_EPLB_DEFINITIONS) + +if(DYNAMIC_EPLB_ARCHS) + list(APPEND DYNAMIC_EPLB_SOURCES moeRebalanceHaloQ.cu) +else() + list(APPEND DYNAMIC_EPLB_SOURCES moeRebalancePlaceholder.cpp) + list(APPEND DYNAMIC_EPLB_DEFINITIONS TRTLLM_DYNAMIC_EPLB_PLACEHOLDER_HALO_Q) + message( + STATUS + "Dynamic EPLB HALO-Q disabled: none of SM100, SM103, or SM107 was requested" + ) +endif() + +if(DYNAMIC_EPLB_ARCHS AND CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL + 13.1) + list(APPEND DYNAMIC_EPLB_SOURCES moeRebalanceTma.cu) +else() + list(APPEND DYNAMIC_EPLB_SOURCES moeRebalancePlaceholder.cpp) + list(APPEND DYNAMIC_EPLB_DEFINITIONS TRTLLM_DYNAMIC_EPLB_PLACEHOLDER_TMA) + if(CMAKE_CUDA_COMPILER_VERSION VERSION_LESS 13.1) + message( + STATUS + "Dynamic EPLB TMA disabled: CUDA 13.1 or newer is required (found ${CMAKE_CUDA_COMPILER_VERSION})" + ) + else() + message( + STATUS + "Dynamic EPLB TMA disabled: none of SM100, SM103, or SM107 was requested" + ) + endif() +endif() + +list(REMOVE_DUPLICATES DYNAMIC_EPLB_SOURCES) +add_library(moe_dynamic_eplb_src OBJECT ${DYNAMIC_EPLB_SOURCES}) + +if(DYNAMIC_EPLB_ARCHS) + # These kernels use SM100-family PTX. The family target avoids inheriting the + # generic kernels_src SM80/SM90 compilation passes. + set_cuda_architectures(moe_dynamic_eplb_src 100f) +endif() + +if(DYNAMIC_EPLB_DEFINITIONS) + target_compile_definitions(moe_dynamic_eplb_src + PRIVATE ${DYNAMIC_EPLB_DEFINITIONS}) +endif() + +set_property(TARGET moe_dynamic_eplb_src PROPERTY POSITION_INDEPENDENT_CODE ON) +set_property(TARGET moe_dynamic_eplb_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS + ON) diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.cu b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.cu new file mode 100644 index 000000000000..66d096e79813 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.cu @@ -0,0 +1,2093 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h" +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlanPolicy.h" + +#include +#include +#include + +#include +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ +constexpr int kMaxEp = 32; +constexpr int kMaxExperts = 384; +constexpr int kMaxBroadcasts = kMaxExperts; +constexpr int kHelperMaskWords = (kMaxBroadcasts + 31) / 32; +constexpr int kMinGroup = 4; +constexpr int kSymPayloadOffset = 64; +using Epoch = std::uint32_t; + +enum PlanField +{ + kPlanExpert = 0, + kPlanLevel = 1, + kPlanGroupBegin = 2, + kPlanGroupSize = 3, + kPlanOwner = 4, + kPlanHelper = 5, +}; + +// Adapt placement storage to the copy layer without exposing scheduler +// internals. +__device__ __forceinline__ sami_copy::PlacementView copy_placement_view(int const* plan) +{ + return {plan[1], kMaxBroadcasts, plan + 2 + kPlanExpert * kMaxBroadcasts, plan + 2 + kPlanLevel * kMaxBroadcasts, + plan + 2 + kPlanGroupBegin * kMaxBroadcasts, plan + 2 + kPlanGroupSize * kMaxBroadcasts, + plan + 2 + kPlanOwner * kMaxBroadcasts, plan + 2 + kPlanHelper * kMaxBroadcasts}; +} + +__device__ __forceinline__ int imin(int a, int b) +{ + return a < b ? a : b; +} + +__device__ __forceinline__ int imax(int a, int b) +{ + return a > b ? a : b; +} + +__device__ __forceinline__ Epoch fetch_add_acq_rel_device(int* address, Epoch value) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + return atom.fetch_add(value, cuda::memory_order_acq_rel); +} + +__device__ __forceinline__ Epoch load_acquire_device(int* address) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + return atom.load(cuda::memory_order_acquire); +} + +__device__ __forceinline__ Epoch load_relaxed_device(int* address) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + return atom.load(cuda::memory_order_relaxed); +} + +__device__ __forceinline__ void store_relaxed_device(int* address, Epoch value) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + atom.store(value, cuda::memory_order_relaxed); +} + +__device__ __forceinline__ void store_release_device(int* address, Epoch value) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + atom.store(value, cuda::memory_order_release); +} + +__device__ __forceinline__ Epoch load_acquire_system(int* address) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + return atom.load(cuda::memory_order_acquire); +} + +__device__ __forceinline__ Epoch load_relaxed_system(int* address) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + return atom.load(cuda::memory_order_relaxed); +} + +__device__ __forceinline__ void store_release_system(int* address, Epoch value) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + atom.store(value, cuda::memory_order_release); +} + +__device__ __forceinline__ bool epoch_before(Epoch observed, Epoch expected) +{ + // RFC-1982-style serial comparison: all arithmetic remains unsigned and + // is unambiguous while peers differ by fewer than half the epoch space. + return observed != expected && expected - observed < (Epoch{1} << 31); +} + +__device__ __forceinline__ void record_status(int* address, Epoch epoch) +{ + // Zero is reserved for "never failed". Preserve a nonzero sticky error at + // the one modular epoch whose bit pattern is zero. + Epoch const value = epoch == 0 ? ~Epoch{0} : epoch; + atomicExch(reinterpret_cast(address), value); +} + +__device__ __forceinline__ int* plan_field(int* plan, int field) +{ + return plan + 2 + field * kMaxBroadcasts; +} + +__device__ __forceinline__ bool group_contains(int rank, int begin, int size) +{ + return rank >= begin && rank < begin + size; +} + +__host__ __device__ __forceinline__ unsigned rank_group_mask(int begin, int size) +{ + return size == 32 ? 0xffffffffu : ((1u << size) - 1u) << begin; +} + +__host__ __device__ __forceinline__ int hierarchy_depth(int ep) +{ + int depth = 1; + for (int level = 1; (1 << level) <= ep; ++level) + { + int groups = 1 << level; + if (ep % groups || ep / groups < kMinGroup) + break; + ++depth; + } + return depth; +} + +/* Distinct experts and total receive capacity independently bound broadcasts. + * Actual S is retained in the public outputs; only scratch uses this bound. */ +__host__ __device__ __forceinline__ int broadcast_capacity(int ep, int experts, int helpers) +{ + int const bounded_helpers = helpers < experts ? helpers : experts; + /* Every hierarchy level divides ep exactly. This cancels two integer + * divisions without changing capacity, including non-power-of-two EP. */ + int const groups = 1 << (hierarchy_depth(ep) - 1); + int const receive_bound = groups * bounded_helpers; + return receive_bound < experts ? receive_bound : experts; +} + +__device__ bool grid_rendezvous( + int* sync, int blocks, unsigned long long spin_cycles, int* status, Epoch* completedEpoch) +{ + __shared__ Epoch generation; + __shared__ int success; + /* The block barrier orders every writer's stores before thread zero's + * device-scope release advertises this CTA. */ + __syncthreads(); + if (threadIdx.x == 0) + { + success = 1; + Epoch const current = load_relaxed_device(sync + 1); + Epoch const target = current + 1; + generation = target; + Epoch const old = fetch_add_acq_rel_device(sync, 1); + if (old == static_cast(blocks - 1)) + { + // The last CTA resets arrivals before publishing the next modular + // epoch. Stream ordering prevents a new launch from entering until + // every CTA in this launch has observed the release. + store_relaxed_device(sync, 0); + store_release_device(sync + 1, target); + } + else if (old >= static_cast(blocks)) + { + record_status(status + 1, target); + success = 0; + } + else + { + unsigned long long const start = clock64(); + Epoch observed = load_relaxed_device(sync + 1); + while (epoch_before(observed, target)) + { + if (clock64() - start > spin_cycles) + { + record_status(status + 1, target); + success = 0; + break; + } + observed = load_relaxed_device(sync + 1); + } + if (!epoch_before(observed, target)) + (void) load_acquire_device(sync + 1); + } + } + __syncthreads(); + *completedEpoch = generation; + return success != 0; +} + +__device__ void local_histogram_pass(int const* routes, int* partial, int route_count, int expert_count, int ctas) +{ + extern __shared__ int shared[]; + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const warps = blockDim.x >> 5; + int const bins = expert_count + 1; + int const shards = warps > 8 ? 8 : warps; + int const shared_count = shards * bins; + for (int index = tid; index < shared_count; index += blockDim.x) + shared[index] = 0; + __syncthreads(); + + int const cta_span = (route_count + ctas - 1) / ctas; + int const cta_begin = blockIdx.x * cta_span; + int const cta_end = imin(route_count, cta_begin + cta_span); + int const warp_span = ((cta_span + warps * 32 - 1) / (warps * 32)) * 32; + int const begin = cta_begin + warp * warp_span; + int const end = imin(cta_end, begin + warp_span); + /* Issue a small load group before consuming it to expose memory-level + * parallelism. The binning is unchanged, and atomicAdd regrouping is + * safe because integer addition commutes. */ + constexpr int kBatch = 4; + for (int base_index = begin; base_index < end; base_index += 32 * kBatch) + { + int expert[kBatch]; +#pragma unroll + for (int step = 0; step < kBatch; ++step) + { + int const index = base_index + step * 32 + lane; + expert[step] = index < end ? routes[index] : -1; + } +#pragma unroll + for (int step = 0; step < kBatch; ++step) + { + int const index = base_index + step * 32 + lane; + if (index < end) + { + int const value = expert[step]; + int const bin = (value >= 0 && value < expert_count) ? value : expert_count; + atomicAdd(shared + (warp & 7) * bins + bin, 1); + } + } + } + __syncthreads(); + + for (int expert = tid; expert < expert_count; expert += blockDim.x) + { + int count = 0; + for (int shard = 0; shard < shards; ++shard) + count += shared[shard * bins + expert]; + partial[blockIdx.x * expert_count + expert] = count; + } +} + +/* Re-rank all local routes on the C-1 worker CTAs while CTA0 performs the + * cross-rank exchange and planner. route_aux[0:2] is a reusable worker + * rendezvous; the following NWORK*(E+1) entries hold per-worker totals. + * Per-warp cursors remain in this CTA's shared memory until materialization. */ +__device__ bool stable_worker_ordinal_pass(int const* routes, int* ordinal, int* route_aux, int* status, + int route_count, int expert_count, int worker, int workers, unsigned long long spin_cycles) +{ + extern __shared__ int shared[]; + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const warps = blockDim.x >> 5; + int const bins = expert_count + 1; + int const shared_count = warps * bins; + for (int index = tid; index < shared_count; index += blockDim.x) + shared[index] = 0; + __syncthreads(); + + int const worker_span = (route_count + workers - 1) / workers; + int const worker_begin = worker * worker_span; + int const worker_end = imin(route_count, worker_begin + worker_span); + int const warp_span = ((worker_span + warps * 32 - 1) / (warps * 32)) * 32; + int const begin = worker_begin + warp * warp_span; + int const end = imin(worker_end, begin + warp_span); + for (int base_index = begin; base_index < end; base_index += 32) + { + int const index = base_index + lane; + bool const active = index < end; + int const expert = active ? routes[index] : -1; + // Padding has its own key so it cannot contribute to real/invalid bins. + int const bin = !active ? bins : ((expert >= 0 && expert < expert_count) ? expert : expert_count); + unsigned const peers = __match_any_sync(0xffffffffu, bin); + int const leader = __ffs(peers) - 1; + int const rank = __popc(peers & ((1u << lane) - 1u)); + int base = 0; + if (active && lane == leader) + base = atomicAdd(shared + warp * bins + bin, __popc(peers)); + // One collective for the entire warp; each lane selects its own leader. + base = __shfl_sync(0xffffffffu, base, leader); + if (active) + ordinal[index] = base + rank; + } + __syncthreads(); + + int* worker_totals = route_aux + 2; + for (int bin = tid; bin < bins; bin += blockDim.x) + { + int total = 0; + for (int local_warp = 0; local_warp < warps; ++local_warp) + total += shared[local_warp * bins + bin]; + worker_totals[worker * bins + bin] = total; + } + + /* The first all-CTA rendezvous guarantees every worker is resident, so a + * worker-only global barrier cannot wait on an unscheduled CTA. */ + Epoch workerEpoch; + if (!grid_rendezvous(route_aux, workers, spin_cycles, status, &workerEpoch)) + return false; + + for (int bin = tid; bin < bins; bin += blockDim.x) + { + int cursor = 0; + for (int previous = 0; previous < worker; ++previous) + cursor += worker_totals[previous * bins + bin]; + for (int local_warp = 0; local_warp < warps; ++local_warp) + { + int index = local_warp * bins + bin; + int count = shared[index]; + shared[index] = cursor; + cursor += count; + } + } + __syncthreads(); + return true; +} + +template +__device__ bool exchange_counts(unsigned long long const* peer_bases, int const* partial, int* shared, int ep, + int expert_count, int local_rank, int ctas, Epoch epoch, unsigned long long spin_cycles, int* status) +{ + int const exchange_ep = FixedEp > 0 ? FixedEp : ep; + int const tid = threadIdx.x; + __shared__ int wait_failed; + if (tid == 0) + wait_failed = 0; + /* Keep the simple accumulation loop; the compiler provides the required + * pipelining without explicit batching. */ + for (int expert = tid; expert < expert_count; expert += blockDim.x) + { + int count = 0; + for (int cta = 0; cta < ctas; ++cta) + count += partial[cta * expert_count + expert]; + shared[expert] = count; + } + __syncthreads(); + + int const parity = epoch & 1; + /* One warp owns each destination row. Its warp barrier orders every + * lane's stores before lane zero's system-scope release publishes the + * complete row to the peer acquire. */ + int const warp = tid >> 5; + int const lane = tid & 31; + for (int destination = warp; destination < exchange_ep; destination += blockDim.x >> 5) + { + int* peer = reinterpret_cast(peer_bases[destination]); + int* row = peer + kSymPayloadOffset + (parity * exchange_ep + local_rank) * expert_count; + /* Keep scalar remote stores for the publication path; the local mirror + * copy below uses vector transfers when alignment permits. */ + for (int expert = lane; expert < expert_count; expert += 32) + row[expert] = shared[expert]; + /* Warp synchronization orders every lane's row stores before lane + * zero resumes; its system-scope release then publishes that complete + * transitive happens-before set to the peer acquire. */ + __syncwarp(); + if (lane == 0) + store_release_system(peer + local_rank, epoch); + } + __syncthreads(); + + int* local = reinterpret_cast(peer_bases[local_rank]); + if (tid < exchange_ep) + { + unsigned long long start = clock64(); + Epoch observed = load_relaxed_system(local + tid); + while (epoch_before(observed, epoch)) + { + if (clock64() - start > spin_cycles) + { + record_status(status + 0, epoch); + atomicExch(&wait_failed, 1); + break; + } + observed = load_relaxed_system(local + tid); + } + if (!epoch_before(observed, epoch)) + (void) load_acquire_system(local + tid); + } + __syncthreads(); + if (wait_failed) + return false; + int const* payload = local + kSymPayloadOffset + parity * exchange_ep * expert_count; + int const total = exchange_ep * expert_count; + /* Vectorize the local mirror copy when shape and alignment permit. */ + if ((total & 3) == 0 && (reinterpret_cast(payload) & 15) == 0) + { + int4 const* source = reinterpret_cast(payload); + int4* destination = reinterpret_cast(shared); + int const vectors = total >> 2; + for (int vector = tid; vector < vectors; vector += blockDim.x) + destination[vector] = source[vector]; + } + else + { + for (int index = tid; index < total; index += blockDim.x) + shared[index] = payload[index]; + } + __syncthreads(); + return true; +} + +__device__ void prepare_halo_m_sparse8(int const* counts, int* histogram, int* load, int* candidates, int experts) +{ + int const tid = threadIdx.x; + int const warp = tid >> 5; + int const lane = tid & 31; + if (warp >= 8) + return; + int const home = experts / 8; + int local_index[2]; + int local_mass[2]; + int rank_load = 0; +#pragma unroll + for (int round = 0; round < 2; ++round) + { + int local = lane + 32 * round; + int mass = -1; + if (local < home) + { + mass = 0; + int expert = warp * home + local; +#pragma unroll + for (int source = 0; source < 8; ++source) + mass += counts[source * experts + expert]; + rank_load += mass; + } + local_index[round] = local; + local_mass[round] = mass; + } + rank_load = __reduce_add_sync(0xffffffffu, rank_load); + if (lane == 0) + load[warp] = rank_load; + for (int slot = 0; slot < 4; ++slot) + { + int best_mass = -1; + int best_local = INT_MAX; +#pragma unroll + for (int round = 0; round < 2; ++round) + { + int mass = local_mass[round]; + int local = local_index[round]; + if (mass > best_mass || (mass == best_mass && local < best_local)) + { + best_mass = mass; + best_local = local; + } + } + int warp_mass = __reduce_max_sync(0xffffffffu, best_mass); + int winner = __reduce_min_sync(0xffffffffu, best_mass == warp_mass ? best_local : INT_MAX); + if (lane == 0) + { + int expert = warp * home + winner; + candidates[warp * 4 + slot] = expert; + histogram[expert] = warp_mass; + } +#pragma unroll + for (int round = 0; round < 2; ++round) + if (local_index[round] == winner) + local_mass[round] = -1; + } + asm volatile("bar.sync 4, 256;" ::: "memory"); +} + +/* Build HALO-M's dense histogram, rank loads and per-owner hot lists with one + * warp per rank. The former thread-0 implementation dominated the kernel. */ +template +__device__ void prepare_halo_m(int const* counts, int* histogram, int* load, int* candidates, int* recv_used, + int* send_used, int* next_index, int* frozen, int ep, int experts, int helpers) +{ + int const prepare_ep = FixedEp > 0 ? FixedEp : ep; + int const prepare_helpers = imin(FixedHelpers > 0 ? FixedHelpers : helpers, experts / prepare_ep); + int const tid = threadIdx.x; + int const warp = tid >> 5; + int const lane = tid & 31; + int const home = experts / prepare_ep; + for (int expert = tid; expert < experts; expert += blockDim.x) + { + int mass = 0; +#pragma unroll + for (int rank = 0; rank < prepare_ep; ++rank) + mass += counts[rank * experts + expert]; + histogram[expert] = mass; + } + __syncthreads(); + + for (int owner = warp; owner < prepare_ep; owner += blockDim.x >> 5) + { + int rank_load = 0; + for (int local = lane; local < home; local += 32) + rank_load += histogram[owner * home + local]; + rank_load = __reduce_add_sync(0xffffffffu, rank_load); + if (lane == 0) + { + load[owner] = rank_load; + recv_used[owner] = 0; + send_used[owner] = 0; + next_index[owner] = 0; + frozen[owner] = 0; + } + + /* EP2 is the largest home: 192 experts, six entries per lane. */ + constexpr int rounds = FixedEp == 8 ? 2 : 6; + int local_index[rounds]; + int local_mass[rounds]; +#pragma unroll + for (int round = 0; round < rounds; ++round) + { + int local = lane + 32 * round; + local_index[round] = local; + local_mass[round] = local < home ? histogram[owner * home + local] : -1; + } + for (int slot = 0; slot < prepare_helpers; ++slot) + { + int best_mass = -1; + int best_local = INT_MAX; +#pragma unroll + for (int round = 0; round < rounds; ++round) + { + int mass = local_mass[round]; + int local = local_index[round]; + if (mass > best_mass || (mass == best_mass && local < best_local)) + { + best_mass = mass; + best_local = local; + } + } + int warp_mass = __reduce_max_sync(0xffffffffu, best_mass); + int winner = __reduce_min_sync(0xffffffffu, best_mass == warp_mass ? best_local : INT_MAX); + if (lane == 0) + candidates[owner * prepare_helpers + slot] = owner * home + winner; +#pragma unroll + for (int round = 0; round < rounds; ++round) + if (local_index[round] == winner) + local_mass[round] = -1; + } + } + __syncthreads(); +} + +/* EP<=8 fast path: every participating lane owns one rank's complete + * mutable HALO-M state. Group and destination decisions use only shuffles, + * avoiding repeated shared-memory rank-state traffic and warp barriers. */ +template +__device__ __forceinline__ int plan_halo_m_register8_impl( + int const* histogram, int lane_load, int total, int* quota, int const* candidates, int* plan, int ep, int helpers) +{ + int const plan_ep = FixedEp > 0 ? FixedEp : ep; + int const plan_helpers = FixedHelpers > 0 ? FixedHelpers : helpers; + int const lane = threadIdx.x & 31; + constexpr unsigned mask = 0xffu; + int lane_recv = 0; + int lane_send = 0; + int lane_next = 0; + int lane_frozen = 0; + + int const target = total / plan_ep + (total % plan_ep != 0); + int const depth = hierarchy_depth(plan_ep); + int* p_expert = plan_field(plan, kPlanExpert); + int* p_level = plan_field(plan, kPlanLevel); + int* p_begin = plan_field(plan, kPlanGroupBegin); + int* p_size = plan_field(plan, kPlanGroupSize); + int* p_owner = plan_field(plan, kPlanOwner); + int* p_helper = plan_field(plan, kPlanHelper); + int broadcasts = 0; + + int trips = 0; + for (int guard = 0; guard < plan_ep * (plan_helpers + 2); ++guard) + { + ++trips; + bool const eligible = lane < plan_ep && !lane_frozen && lane_load > target; + int owner, owner_load; + if constexpr (NarrowKeys) + { + unsigned const key = __reduce_max_sync( + mask, eligible ? ((static_cast(lane_load) << 4) | (15u - static_cast(lane))) : 0u); + if (!key) + break; + owner = 15 - static_cast(key & 15u); + owner_load = static_cast(key >> 4); + } + else + { + owner_load = __reduce_max_sync(mask, eligible ? lane_load : -1); + if (owner_load < 0) + break; + owner = __reduce_min_sync(mask, eligible && lane_load == owner_load ? lane : INT_MAX); + } + + unsigned const packed + = __shfl_sync(mask, (static_cast(lane_send) << 4) | static_cast(lane_next), owner, 8); + int const owner_send = static_cast(packed >> 4); + int const owner_next = static_cast(packed & 15u); + if (owner_send >= plan_helpers || owner_next >= plan_helpers) + { + if (lane == owner) + lane_frozen = 1; + continue; + } + + /* owner and owner_next are now warp-uniform, so every lane can read + * candidates[] and histogram[] directly. Both are written by + * prepare_halo_m_sparse8 before its bar.sync and never inside this + * loop, so a uniform load returns exactly what the owner would have + * broadcast. The increment stays owner-only and stays ahead of the + * mass test, as before. */ + int const expert = candidates[owner * plan_helpers + owner_next]; + if (lane == owner) + ++lane_next; + int const mass = histogram[expert]; + if (mass == 0) + { + if (lane == owner) + lane_frozen = 1; + continue; + } + + int give = imin(owner_load - target, mass); + if (give <= 0) + { + if (lane == owner) + lane_frozen = 1; + continue; + } + + int positive_target = imax(target, 1); + int need = imax(1, mass / positive_target + (mass % positive_target != 0)); + int wanted = depth - 1; + for (int level = depth - 1; level >= 0; --level) + { + if (plan_ep / (1 << level) >= need) + { + wanted = level; + break; + } + } + + int chosen_level = -1; + int chosen_begin = -1; + int chosen_size = 0; + if constexpr (FixedEp == 8 && FixedHelpers == 4) + { + int slack = lane != owner ? imax(0, target - lane_load) : 0; + /* Two independent full-mask sums instead of one lane-varying-mask + * reduce followed by two broadcasts. Integer addition is exact, so + * masking the addend gives the same half sums; the two reduces do + * not depend on each other, so the chain is depth 1 instead of 3, + * and the redux.sync no longer runs with a per-lane membermask. */ + int room0 = __reduce_add_sync(mask, lane < 4 ? slack : 0); + int room4 = __reduce_add_sync(mask, lane >= 4 ? slack : 0); + unsigned receive_mask = __ballot_sync(mask, lane_recv < plan_helpers); + bool valid0 = (receive_mask & 0x0fu) == 0x0fu && room0 > 0; + bool valid4 = (receive_mask & 0xf0u) == 0xf0u && room4 > 0; + bool want_half = mass <= 4 * target; + if (want_half && (valid0 || valid4)) + { + chosen_level = 1; + chosen_begin = valid0 && (!valid4 || room0 >= room4) ? 0 : 4; + chosen_size = 4; + } + else if (receive_mask == mask && room0 + room4 > 0) + { + chosen_level = 0; + chosen_begin = 0; + chosen_size = 8; + } + } + else + { + for (int level = wanted; level >= 0 && chosen_level < 0; --level) + { + int size = plan_ep / (1 << level); + bool active = lane < plan_ep; + int begin = (lane / size) * size; + unsigned group_mask = rank_group_mask(begin, size); + int local_room = active && lane != owner ? imax(0, target - lane_load) : 0; + int room = 0; + if (active) + room = __reduce_add_sync(group_mask, local_room); + unsigned receive_mask = __ballot_sync(mask, active && lane_recv < plan_helpers); + bool group_leader = active && lane == begin && (receive_mask & group_mask) == group_mask; + unsigned leader_mask = __ballot_sync(mask, group_leader); + int best_room = -1; + int best_begin = INT_MAX; + if (leader_mask) + { + int const first_leader = __ffs(leader_mask) - 1; + if constexpr (NarrowKeys) + { + unsigned key = 0; + if (group_leader) + key = __reduce_max_sync( + leader_mask, (static_cast(room) << 4) | (15u - static_cast(begin))); + key = __shfl_sync(mask, key, first_leader, 8); + best_room = static_cast(key >> 4); + best_begin = 15 - static_cast(key & 15u); + } + else + { + if (group_leader) + { + best_room = __reduce_max_sync(leader_mask, room); + best_begin = __reduce_min_sync(leader_mask, room == best_room ? begin : INT_MAX); + } + best_room = __shfl_sync(mask, best_room, first_leader, 8); + best_begin = __shfl_sync(mask, best_begin, first_leader, 8); + } + } + if (best_begin != INT_MAX && best_room > 0) + { + chosen_level = level; + chosen_begin = best_begin; + chosen_size = size; + } + } + } + if (chosen_level < 0) + { + if (lane == owner) + lane_frozen = 1; + continue; + } + if (broadcasts >= kMaxBroadcasts) + break; + + int lane_quota = lane == owner ? mass : 0; + int remaining = give; + unsigned visited = 0; + for (int step = 0; step < chosen_size && remaining > 0; ++step) + { + bool can_receive = lane >= chosen_begin && lane < chosen_begin + chosen_size && lane != owner + && !(visited & (1u << lane)); + int destination, destination_load; + if constexpr (NarrowKeys) + { + unsigned const key = __reduce_min_sync(mask, + can_receive ? ((static_cast(lane_load) << 4) | static_cast(lane)) + : 0xffffffffu); + if (key == 0xffffffffu) + break; + destination = static_cast(key & 15u); + destination_load = static_cast(key >> 4); + } + else + { + destination_load = __reduce_min_sync(mask, can_receive ? lane_load : INT_MAX); + destination = __reduce_min_sync(mask, can_receive && lane_load == destination_load ? lane : INT_MAX); + if (destination == INT_MAX) + break; + } + visited |= 1u << destination; + int room = target - destination_load; + if (room <= 0) + break; + int amount = imin(room, remaining); + if (lane == owner) + lane_quota -= amount; + if (lane == destination) + lane_quota += amount; + remaining -= amount; + } + + int moved = give - remaining; + if (moved <= 0) + { + if (lane == owner) + lane_frozen = 1; + continue; + } + + int b = broadcasts++; + if (lane == 0) + { + p_expert[b] = expert; + p_level[b] = chosen_level; + p_begin[b] = chosen_begin; + p_size[b] = chosen_size; + p_owner[b] = owner; + p_helper[b] = -1; + } + if (lane < plan_ep) + quota[b * plan_ep + lane] = lane_quota; + + bool in_group = lane >= chosen_begin && lane < chosen_begin + chosen_size; + if (in_group) + ++lane_recv; + if (lane == owner) + { + ++lane_send; + lane_load -= moved; + } + else if (in_group) + { + lane_load += lane_quota; + } + } + return broadcasts; +} + +/* Select the key representation once, from the actual complete rank load. + * This guard keeps owner/destination loads and group-room keys representable. + * Keep the wide path for large valid int32 inputs without putting a runtime + * representation branch on each iteration of either planning loop. */ +template +__device__ __forceinline__ int plan_halo_m_register8( + int const* histogram, int const* load, int* quota, int const* candidates, int* plan, int ep, int helpers) +{ + int const plan_ep = FixedEp > 0 ? FixedEp : ep; + int const lane = threadIdx.x & 31; + int const lane_load = lane < plan_ep ? load[lane] : 0; + int const total = __reduce_add_sync(0xffu, lane_load); + if (total <= 0x0fffffff) + return plan_halo_m_register8_impl( + histogram, lane_load, total, quota, candidates, plan, ep, helpers); + return plan_halo_m_register8_impl( + histogram, lane_load, total, quota, candidates, plan, ep, helpers); +} + +/* Exact HALO-M state machine executed by warp 0. Each lane owns one rank; + * warp reductions replace all serial rank/group/destination scans. */ +__device__ int plan_halo_m_warp(int const* histogram, int* load, int* quota, int const* candidates, int* recv_used, + int* send_used, int* next_index, int* frozen, int* plan, int ep, int experts, int helpers) +{ + int const candidate_count = imin(helpers, experts / ep); + int const lane = threadIdx.x & 31; + int const reduce_width = ep <= 8 ? 8 : 32; + int const reduce_start = reduce_width >> 1; + unsigned const reduce_mask = ep <= 8 ? 0xffu : 0xffffffffu; + int total = lane < ep ? load[lane] : 0; + for (int offset = reduce_start; offset; offset >>= 1) + { + int const other = __shfl_down_sync(reduce_mask, total, offset, reduce_width); + if (lane + offset < reduce_width) + total += other; + } + total = __shfl_sync(reduce_mask, total, 0, reduce_width); + int const target = total / ep + (total % ep != 0); + int const depth = hierarchy_depth(ep); + int* p_expert = plan_field(plan, kPlanExpert); + int* p_level = plan_field(plan, kPlanLevel); + int* p_begin = plan_field(plan, kPlanGroupBegin); + int* p_size = plan_field(plan, kPlanGroupSize); + int* p_owner = plan_field(plan, kPlanOwner); + int* p_helper = plan_field(plan, kPlanHelper); + int broadcasts = 0; + + for (int guard = 0; guard < ep * (candidate_count + 2); ++guard) + { + bool eligible = lane < ep && !frozen[lane] && load[lane] > target; + int best_load = eligible ? load[lane] : -1; + int best_rank = eligible ? lane : INT_MAX; + for (int offset = reduce_start; offset; offset >>= 1) + { + int other_load = __shfl_down_sync(reduce_mask, best_load, offset, reduce_width); + int other_rank = __shfl_down_sync(reduce_mask, best_rank, offset, reduce_width); + if (lane + offset < reduce_width + && (other_load > best_load || (other_load == best_load && other_rank < best_rank))) + { + best_load = other_load; + best_rank = other_rank; + } + } + int owner = __shfl_sync(reduce_mask, best_rank, 0, reduce_width); + if (owner == INT_MAX) + break; + /* One lane owns the capacity check and shared cursor transition. + * Publish either the chosen expert or -1 for exhaustion explicitly; + * no reader can observe a partially advanced cursor. */ + int expert = -1; + if (lane == 0) + { + if (send_used[owner] >= helpers || next_index[owner] >= candidate_count) + { + frozen[owner] = 1; + } + else + { + int const candidate_index = next_index[owner]++; + expert = candidates[owner * candidate_count + candidate_index]; + } + } + expert = __shfl_sync(reduce_mask, expert, 0, reduce_width); + __syncwarp(reduce_mask); + if (expert < 0) + continue; + int mass = histogram[expert]; + if (mass == 0) + { + if (lane == 0) + frozen[owner] = 1; + __syncwarp(reduce_mask); + continue; + } + int give = imin(load[owner] - target, mass); + if (give <= 0) + { + if (lane == 0) + frozen[owner] = 1; + __syncwarp(reduce_mask); + continue; + } + + int need = imax(1, mass / imax(target, 1) + (mass % imax(target, 1) != 0)); + int wanted = depth - 1; + for (int level = depth - 1; level >= 0; --level) + { + if (ep / (1 << level) >= need) + { + wanted = level; + break; + } + } + int chosen_level = -1; + int chosen_begin = -1; + int chosen_size = 0; + for (int level = wanted; level >= 0 && chosen_level < 0; --level) + { + int size = ep / (1 << level); + int begin = lane * size; + bool active = begin < ep; + bool affordable = active; + int room = 0; + if (active) + { + for (int rank = begin; rank < begin + size; ++rank) + { + affordable &= recv_used[rank] < helpers; + if (rank != owner) + room += imax(0, target - load[rank]); + } + } + int candidate_room = affordable ? room : -1; + int candidate_begin = affordable ? begin : INT_MAX; + for (int offset = reduce_start; offset; offset >>= 1) + { + int other_room = __shfl_down_sync(reduce_mask, candidate_room, offset, reduce_width); + int other_begin = __shfl_down_sync(reduce_mask, candidate_begin, offset, reduce_width); + if (lane + offset < reduce_width + && (other_room > candidate_room || (other_room == candidate_room && other_begin < candidate_begin))) + { + candidate_room = other_room; + candidate_begin = other_begin; + } + } + int best_room = __shfl_sync(reduce_mask, candidate_room, 0, reduce_width); + int best_begin = __shfl_sync(reduce_mask, candidate_begin, 0, reduce_width); + if (best_begin != INT_MAX && best_room > 0) + { + chosen_level = level; + chosen_begin = best_begin; + chosen_size = size; + } + } + if (chosen_level < 0) + { + if (lane == 0) + frozen[owner] = 1; + __syncwarp(reduce_mask); + continue; + } + if (broadcasts >= kMaxBroadcasts) + break; + + int b = broadcasts++; + if (lane == 0) + { + p_expert[b] = expert; + p_level[b] = chosen_level; + p_begin[b] = chosen_begin; + p_size[b] = chosen_size; + p_owner[b] = owner; + p_helper[b] = -1; + } + if (lane < ep) + quota[b * ep + lane] = 0; + __syncwarp(reduce_mask); + if (lane == 0) + quota[b * ep + owner] = mass; + __syncwarp(reduce_mask); + + int remaining = give; + unsigned visited = 0; + for (int step = 0; step < chosen_size && remaining > 0; ++step) + { + bool can_receive = lane >= chosen_begin && lane < chosen_begin + chosen_size && lane != owner + && !(visited & (1u << lane)); + int cold_load = can_receive ? load[lane] : INT_MAX; + int cold_rank = can_receive ? lane : INT_MAX; + for (int offset = reduce_start; offset; offset >>= 1) + { + int other_load = __shfl_down_sync(reduce_mask, cold_load, offset, reduce_width); + int other_rank = __shfl_down_sync(reduce_mask, cold_rank, offset, reduce_width); + if (lane + offset < reduce_width + && (other_load < cold_load || (other_load == cold_load && other_rank < cold_rank))) + { + cold_load = other_load; + cold_rank = other_rank; + } + } + int destination = __shfl_sync(reduce_mask, cold_rank, 0, reduce_width); + if (destination == INT_MAX) + break; + visited |= 1u << destination; + int room = target - load[destination]; + if (room <= 0) + continue; + int amount = imin(room, remaining); + if (lane == 0) + { + quota[b * ep + owner] -= amount; + quota[b * ep + destination] += amount; + } + remaining -= amount; + __syncwarp(reduce_mask); + } + int moved = give - remaining; + if (moved <= 0) + { + --broadcasts; + if (lane == 0) + frozen[owner] = 1; + __syncwarp(reduce_mask); + continue; + } + if (lane >= chosen_begin && lane < chosen_begin + chosen_size) + ++recv_used[lane]; + if (lane == owner) + { + ++send_used[owner]; + load[owner] -= moved; + } + else if (lane >= chosen_begin && lane < chosen_begin + chosen_size) + { + load[lane] += quota[b * ep + lane]; + } + __syncwarp(reduce_mask); + } + return broadcasts; +} + +__device__ __forceinline__ int warp_sum(int value) +{ + for (int offset = 16; offset; offset >>= 1) + { + int const other = __shfl_down_sync(0xffffffffu, value, offset); + if ((threadIdx.x & 31) + offset < 32) + value += other; + } + return __shfl_sync(0xffffffffu, value, 0); +} + +__device__ __forceinline__ int warp_exclusive_scan(int value) +{ + int sum = value; + int const lane = threadIdx.x & 31; + for (int offset = 1; offset < 32; offset <<= 1) + { + int other = __shfl_up_sync(0xffffffffu, sum, offset); + if (lane >= offset) + sum += other; + } + return sum - value; +} + +__device__ void halo_q_repair(int receiver, int donor, int broadcasts, int const* order, int const* host_mask, + int const* counts, int const* p_expert, int* quota, int ep, int experts) +{ + int const lane = threadIdx.x & 31; + if (broadcasts > 32) + { + int local_pull = 0, local_give = 0; + for (int base = 0; base < broadcasts; base += 32) + { + int const b = base + lane < broadcasts ? order[base + lane] : -1; + if (b >= 0 && (host_mask[b] & (1u << receiver)) && (host_mask[b] & (1u << donor))) + { + int const expert = p_expert[b]; + int const cr = counts[receiver * experts + expert]; + int const cd = counts[donor * experts + expert]; + int const qr = quota[b * ep + receiver]; + int const qd = quota[b * ep + donor]; + local_pull += imin(imax(0, cr - qr), imax(0, qd - cd)); + local_give += imax(0, qr - cr); + } + } + int const delta = imin(warp_sum(local_pull), warp_sum(local_give)); + int pull_before_tiles = 0, give_before_tiles = 0; + for (int base = 0; base < broadcasts; base += 32) + { + int const b = base + lane < broadcasts ? order[base + lane] : -1; + int pull = 0, give = 0; + if (b >= 0 && (host_mask[b] & (1u << receiver)) && (host_mask[b] & (1u << donor))) + { + int const expert = p_expert[b]; + int const cr = counts[receiver * experts + expert]; + int const cd = counts[donor * experts + expert]; + int const qr = quota[b * ep + receiver]; + int const qd = quota[b * ep + donor]; + pull = imin(imax(0, cr - qr), imax(0, qd - cd)); + give = imax(0, qr - cr); + } + int const pull_before = pull_before_tiles + warp_exclusive_scan(pull); + int const give_before = give_before_tiles + warp_exclusive_scan(give); + int const take = imin(pull, imax(0, delta - pull_before)); + int const release = imin(give, imax(0, delta - give_before)); + pull_before_tiles += warp_sum(pull); + give_before_tiles += warp_sum(give); + if (b >= 0) + { + quota[b * ep + receiver] += take - release; + quota[b * ep + donor] += release - take; + } + __syncwarp(); + } + return; + } + int b = lane < broadcasts ? order[lane] : -1; + int pull = 0; + int give = 0; + if (b >= 0 && (host_mask[b] & (1u << receiver)) && (host_mask[b] & (1u << donor))) + { + int expert = p_expert[b]; + int cr = counts[receiver * experts + expert]; + int cd = counts[donor * experts + expert]; + int qr = quota[b * ep + receiver]; + int qd = quota[b * ep + donor]; + pull = imin(imax(0, cr - qr), imax(0, qd - cd)); + give = imax(0, qr - cr); + } + int delta = imin(warp_sum(pull), warp_sum(give)); + int pull_before = warp_exclusive_scan(pull); + int give_before = warp_exclusive_scan(give); + int pull_take = imin(pull, imax(0, delta - pull_before)); + int give_take = imin(give, imax(0, delta - give_before)); + if (b >= 0) + { + quota[b * ep + receiver] += pull_take - give_take; + quota[b * ep + donor] += give_take - pull_take; + } + __syncwarp(); +} + +__device__ __forceinline__ int warp_sum8(int value) +{ + return __reduce_add_sync(0xffu, value); +} + +__device__ __forceinline__ int warp_exclusive_scan8(int value) +{ + int sum = value; + int const lane = threadIdx.x & 7; + for (int offset = 1; offset < 8; offset <<= 1) + { + int other = __shfl_up_sync(0xffu, sum, offset, 8); + if (lane >= offset) + sum += other; + } + return sum - value; +} + +__device__ __forceinline__ void halo_q_direction8( + bool common, int receiver_count, int donor_count, int& receiver_quota, int& donor_quota) +{ + int pull = common ? imin(imax(0, receiver_count - receiver_quota), imax(0, donor_quota - donor_count)) : 0; + int give = common ? imax(0, receiver_quota - receiver_count) : 0; + int pull_total = warp_sum8(pull); + int give_total = warp_sum8(give); + int delta = imin(pull_total, give_total); + if (!delta) + return; + int pull_before = warp_exclusive_scan8(pull); + int give_before = warp_exclusive_scan8(give); + int pull_take = imin(pull, imax(0, delta - pull_before)); + int give_take = imin(give, imax(0, delta - give_before)); + receiver_quota += pull_take - give_take; + donor_quota += give_take - pull_take; +} + +/* Fuse both ordered directions of one rank pair. Each broadcast lane loads + * its two count/quota entries once, executes left<-right then right<-left in + * registers, and writes the final pair once. + * + * `b`, `hosts` and `expert` are the lane's own broadcast identity and do not + * change between phases or rounds, so the caller hoists them. Reading them + * here put a four-deep dependent load chain (order[lane] -> b -> host_mask[b] + * and the GLOBAL p_expert[b] -> expert -> counts[...]) at the head of all + * fourteen phases. */ +__device__ __forceinline__ void halo_q_pair8( + int left, int right, int b, int hosts, int expert, int const* counts, int* quota, int ep, int experts) +{ + int const lane = threadIdx.x & 31; + if (lane >= 8) + return; + bool common = false; + int count_left = 0; + int count_right = 0; + int quota_left = 0; + int quota_right = 0; + if (b >= 0) + { + common = (hosts & (1u << left)) && (hosts & (1u << right)); + if (common) + { + count_left = counts[left * experts + expert]; + count_right = counts[right * experts + expert]; + quota_left = quota[b * ep + left]; + quota_right = quota[b * ep + right]; + } + } + + halo_q_direction8(common, count_left, count_right, quota_left, quota_right); + halo_q_direction8(common, count_right, count_left, quota_right, quota_left); + + if (common) + { + quota[b * ep + left] = quota_left; + quota[b * ep + right] = quota_right; + } +} + +__device__ __forceinline__ void halo_q_active_barrier(int threads) +{ + asm volatile("bar.sync 1, %0;" ::"r"(threads) : "memory"); +} + +__device__ __forceinline__ void scheduler_tail_barrier() +{ + asm volatile("bar.sync 3, 64;" ::: "memory"); +} + +/* Build order[] and host_mask[] with one lane per broadcast. + * + * This replaces serial insertion-sort and mask-construction loops while + * preserving the same stable order. + * + * order[]: lane b's position is the number of broadcasts that sort before it. + * The `j < lane` tiebreak makes this the same stable order the + * insertion sort produced (its shift test is a strict `>`), so the + * result is bit-identical, not merely equivalent. + * host_mask: the destination group is a contiguous rank range, so the mask is + * closed form; no loop over group members is needed. + * rank_group_mask handles the full EP32 word without a shift by32. + */ +__device__ __forceinline__ void halo_q_build_order_and_masks(int const* p_expert, int const* p_begin, int const* p_size, + int const* p_owner, int* order, int* host_mask, int broadcasts) +{ + int const lane = threadIdx.x & 31; + if (broadcasts > 32) + { + for (int b = lane; b < broadcasts; b += 32) + { + int const mine = p_expert[b]; + int position = 0; + for (int j = 0; j < broadcasts; ++j) + position += p_expert[j] < mine || (p_expert[j] == mine && j < b); + order[position] = b; + host_mask[b] = static_cast(rank_group_mask(p_begin[b], p_size[b]) | (1u << p_owner[b])); + } + return; + } + /* Read p_expert cooperatively and distribute values with shuffles. All + * lanes must participate before inactive lanes return; inactive lanes + * carry INT_MAX and never contribute to a real comparison. */ + bool const active = lane < broadcasts; + int const mine = active ? p_expert[lane] : INT_MAX; + int position = 0; + for (int j = 0; j < broadcasts; ++j) + { + int const other = __shfl_sync(0xffffffffu, mine, j); + position += (other < mine) || (other == mine && j < lane); + } + if (!active) + return; + order[position] = lane; + host_mask[lane] = static_cast(rank_group_mask(p_begin[lane], p_size[lane]) | (1u << p_owner[lane])); +} + +template +__device__ void run_halo_q(int const* counts, int* quota, int* plan, int broadcasts, int ep, int experts, int* ring, + int* order, int* host_mask) +{ + int const plan_ep = FixedEp > 0 ? FixedEp : ep; + int const tid = threadIdx.x; + int const warp = tid >> 5; + int const lane = tid & 31; + int* p_expert = plan_field(plan, kPlanExpert); + int* p_begin = plan_field(plan, kPlanGroupBegin); + int* p_size = plan_field(plan, kPlanGroupSize); + int* p_owner = plan_field(plan, kPlanOwner); + if (plan_ep <= 8 && broadcasts <= 8) + { + if (warp < plan_ep / 2) + { + if (warp == 0) + halo_q_build_order_and_masks(p_expert, p_begin, p_size, p_owner, order, host_mask, broadcasts); + halo_q_active_barrier((plan_ep / 2) * 32); + /* Loop-invariant broadcast identity, read once after the barrier + * that publishes order[] and host_mask[]. All four participating + * warps read the same lane-indexed entries. */ + int const my_b = lane < broadcasts ? order[lane] : -1; + int const my_hosts = my_b >= 0 ? host_mask[my_b] : 0; + int const my_expert = my_b >= 0 ? p_expert[my_b] : -1; + if (broadcasts >= 2) + { + int const ring_size = plan_ep - 1; + /* Use two unconditional repair rounds. Phases with no required repair + * are no-ops, so this preserves the conditional algorithm's result. */ + int const repair_rounds = 2; + for (int round = 0; round < repair_rounds; ++round) + { + for (int phase = 0; phase < plan_ep - 1; ++phase) + { + int right_index = warp == 0 ? 0 : plan_ep - 1 - warp; + int right_position = right_index - phase; + if (right_position < 0) + right_position += ring_size; + int right = 1 + right_position; + int left = 0; + if (warp != 0) + { + int left_position = warp - phase; + if (left_position < 0) + left_position += ring_size; + left = 1 + left_position; + } + halo_q_pair8(left, right, my_b, my_hosts, my_expert, counts, quota, plan_ep, experts); + if constexpr (FixedEp == 8) + { + /* Must stay unconditional. Omitting it after a + * round's last phase races twice: warp 0 enters the + * next round's phase 0 while warp 3 still writes the + * same quota[], and after the final phase warps 1-3 + * leave run_halo_q while warp 0 -- which only meets + * warp 8 at scheduler_tail_barrier -- walks into + * build_route_prefix and reads every quota[]. */ + halo_q_active_barrier(128); + } + else + { + halo_q_active_barrier((plan_ep / 2) * 32); + } + } + } + } + } + /* The unconditional repair rounds above replace the conditional scan. */ + return; + } + if (warp == 0) + halo_q_build_order_and_masks(p_expert, p_begin, p_size, p_owner, order, host_mask, broadcasts); + if (tid < plan_ep - 1) + ring[tid] = tid + 1; + __syncthreads(); + if (broadcasts < 2) + return; + + for (int round = 0; round < 2; ++round) + { + for (int phase = 0; phase < plan_ep - 1; ++phase) + { + if (warp < plan_ep / 2) + { + int left, right; + if (warp == 0) + { + left = 0; + right = ring[0]; + } + else + { + left = ring[warp]; + right = ring[plan_ep - 1 - warp]; + } + halo_q_repair(left, right, broadcasts, order, host_mask, counts, p_expert, quota, plan_ep, experts); + halo_q_repair(right, left, broadcasts, order, host_mask, counts, p_expert, quota, plan_ep, experts); + } + __syncthreads(); + if (tid == 0) + { + int last = ring[plan_ep - 2]; + for (int index = plan_ep - 2; index > 0; --index) + ring[index] = ring[index - 1]; + ring[0] = last; + } + __syncthreads(); + } + } + (void) lane; +} + +__device__ void color_and_publish_local(int* plan, int broadcasts, int ep, int helpers, int local_rank, int* used, + int* out_ids, int* out_levels, int* out_owners, int* status, Epoch epoch) +{ + int* p_expert = plan_field(plan, kPlanExpert); + int* p_level = plan_field(plan, kPlanLevel); + int* p_begin = plan_field(plan, kPlanGroupBegin); + int* p_size = plan_field(plan, kPlanGroupSize); + int* p_owner = plan_field(plan, kPlanOwner); + int* p_helper = plan_field(plan, kPlanHelper); + int const active_helpers = imin(helpers, kMaxBroadcasts); + int const words = (active_helpers + 31) / 32; + for (int rank = 0; rank < ep; ++rank) + for (int word = 0; word < words; ++word) + used[rank * kHelperMaskWords + word] = 0; + for (int level = 0; level < hierarchy_depth(ep); ++level) + { + for (int b = 0; b < broadcasts; ++b) + { + if (p_level[b] != level) + continue; + unsigned unavailable[kHelperMaskWords] = {0}; + for (int rank = p_begin[b]; rank < p_begin[b] + p_size[b]; ++rank) + for (int word = 0; word < words; ++word) + unavailable[word] |= static_cast(used[rank * kHelperMaskWords + word]); + int helper = -1; + for (int slot = 0; slot < active_helpers; ++slot) + if (!(unavailable[slot / 32] & (1u << (slot % 32)))) + { + helper = slot; + break; + } + if (helper < 0) + { + record_status(status + 4, epoch); + helper = 0; + } + p_helper[b] = helper; + for (int rank = p_begin[b]; rank < p_begin[b] + p_size[b]; ++rank) + used[rank * kHelperMaskWords + helper / 32] |= static_cast(1u << (helper % 32)); + } + } + for (int slot = 0; slot < helpers; ++slot) + { + out_ids[slot] = -1; + out_levels[slot] = -1; + out_owners[slot] = -1; + } + for (int b = 0; b < broadcasts; ++b) + { + if (!group_contains(local_rank, p_begin[b], p_size[b])) + continue; + int helper = p_helper[b]; + if (out_ids[helper] != -1) + record_status(status + 4, epoch); + out_ids[helper] = p_expert[b]; + out_levels[helper] = p_level[b]; + out_owners[helper] = p_owner[b]; + } +} + +__device__ void build_route_prefix(int const* counts, int const* quota, int* prefix, int const* plan, int broadcasts, + int ep, int experts, int local_rank, int participating_threads) +{ + int* p_expert = plan_field(const_cast(plan), kPlanExpert); + for (int b = threadIdx.x; b < broadcasts; b += participating_threads) + { + int expert = p_expert[b]; + int source_before = 0; + for (int rank = 0; rank < local_rank; ++rank) + { + int diagonal = imin(counts[rank * experts + expert], quota[b * ep + rank]); + source_before += counts[rank * experts + expert] - diagonal; + } + int source_count = counts[local_rank * experts + expert]; + int source_diagonal = imin(source_count, quota[b * ep + local_rank]); + int residual_supply = source_count - source_diagonal; + /* The loop used to run to destination == ep. route_destination only + * ever loads prefix[destination + step] under `candidate < ep`, so the + * ep entry is never read, and the accumulation at ep-1 fed nothing but + * that dead entry. Both are gone; entry 0 is still written because it + * costs one store and leaving it stale would depend on nobody ever + * reading it, which is a weaker guarantee than writing the right value. */ + int destination_before = 0; + for (int destination = 0; destination < ep; ++destination) + { + int overlap = destination_before - source_before; + overlap = imax(0, imin(residual_supply, overlap)); + int diagonal_prefix = local_rank < destination ? source_diagonal : 0; + prefix[b * (ep + 1) + destination] = diagonal_prefix + overlap; + if (destination + 1 < ep) + { + int diagonal = imin(counts[destination * experts + expert], quota[b * ep + destination]); + destination_before += quota[b * ep + destination] - diagonal; + } + } + } +} + +__device__ __forceinline__ int* worker_expert_lut(int experts) +{ + extern __shared__ int shared[]; + return shared + (blockDim.x >> 5) * (experts + 1); +} + +__device__ void prepare_worker_default_slots(int* expert_lut, int ep, int experts, int helpers) +{ + int const home = experts / ep; + int const local_slots = home + helpers; + for (int expert = threadIdx.x; expert < experts; expert += blockDim.x) + { + int owner = expert / home; + int local = expert - owner * home; + /* Negative values distinguish the precomputed cold slot from a + * nonnegative broadcast index. Slot zero therefore maps to -1. */ + expert_lut[expert] = -(owner * local_slots + local) - 1; + } + __syncthreads(); +} + +__device__ void prepare_worker_selected_experts(int* expert_lut, int const* plan) +{ + int const broadcasts = plan[1]; + int const* p_expert = plan + 2 + kPlanExpert * kMaxBroadcasts; + for (int b = threadIdx.x; b < broadcasts; b += blockDim.x) + expert_lut[p_expert[b]] = b; + __syncthreads(); +} + +__device__ bool wait_for_plan_epoch(int* marker, Epoch epoch, unsigned long long spin_cycles, int* status, int* ready) +{ + if (threadIdx.x == 0) + { + *ready = 1; + unsigned long long start = clock64(); + Epoch observed = load_relaxed_device(marker); + while (epoch_before(observed, epoch)) + { + if (clock64() - start > spin_cycles) + { + record_status(status + 2, epoch); + *ready = 0; + break; + } + observed = load_relaxed_device(marker); + } + if (!epoch_before(observed, epoch)) + (void) load_acquire_device(marker); + } + __syncthreads(); + return *ready != 0; +} + +/* Prefix entries are monotonic. Binary lifting returns the same number of + * thresholds <= ordinal as the old linear scan, including repeated entries. */ +__device__ __forceinline__ int route_destination(int const* prefix, int ordinal, int ep) +{ + int destination = 0; +#pragma unroll + for (int step = 16; step; step >>= 1) + { + int candidate = destination + step; + if (candidate < ep && ordinal >= prefix[candidate]) + destination = candidate; + } + return destination; +} + +/* Once HALO-M publishes its selected experts, workers can finish every cold + * route while CTA0 runs HALO-Q. Selected route indices are compacted into + * disjoint slices of route_aux; their ordinal values remain in output. */ +__device__ void pre_materialize_and_compact_worker_routes(int const* routes, int* output, int* route_aux, int* status, + int route_count, int experts, int worker, int workers, Epoch epoch) +{ + extern __shared__ int shared[]; + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const warps = blockDim.x >> 5; + int const bins = experts + 1; + int* expert_lut = worker_expert_lut(experts); + int* hot_indices = route_aux + 2 + workers * bins; + + int const worker_span = (route_count + workers - 1) / workers; + int const worker_begin = worker * worker_span; + int const worker_end = imin(route_count, worker_begin + worker_span); + int const warp_span = ((worker_span + warps * 32 - 1) / (warps * 32)) * 32; + int const begin = worker_begin + warp * warp_span; + int const end = imin(worker_end, begin + warp_span); + int hot_count = 0; + for (int base = begin; base < end; base += 32) + { + int index = base + lane; + bool active = index < end; + int expert = active ? routes[index] : -1; + int slot = -1; + bool hot = false; + if (active && expert >= 0 && expert < experts) + { + int entry = expert_lut[expert]; + hot = entry >= 0; + if (!hot) + slot = -entry - 1; + } + else if (active && expert != -1) + { + record_status(status + 3, epoch); + } + unsigned hot_mask = __ballot_sync(0xffffffffu, hot); + if (hot) + { + int offset = __popc(hot_mask & ((1u << lane) - 1u)); + hot_indices[begin + hot_count + offset] = index; + } + else if (active) + { + output[index] = slot; + } + hot_count += __popc(hot_mask); + } + if (lane == 0) + shared[warp * bins + experts] = hot_count; +} + +/* The final epoch only has to resolve selected routes. Stable ordinal + * cursors remain in the per-warp shared histogram from the initial pass. */ +__device__ void materialize_worker_hot_routes(int const* routes, int* output, int const* route_prefix, int const* plan, + int* route_aux, int* status, int route_count, int ep, int experts, int helpers, int worker, int workers, + Epoch epoch) +{ + extern __shared__ int shared[]; + int const lane = threadIdx.x & 31; + int const warp = threadIdx.x >> 5; + int const warps = blockDim.x >> 5; + int const bins = experts + 1; + int* expert_lut = worker_expert_lut(experts); + int* hot_indices = route_aux + 2 + workers * bins; + int const home = experts / ep; + int const local_slots = home + helpers; + int const* p_begin = plan + 2 + kPlanGroupBegin * kMaxBroadcasts; + int const* p_size = plan + 2 + kPlanGroupSize * kMaxBroadcasts; + int const* p_owner = plan + 2 + kPlanOwner * kMaxBroadcasts; + int const* p_helper = plan + 2 + kPlanHelper * kMaxBroadcasts; + + int const worker_span = (route_count + workers - 1) / workers; + int const worker_begin = worker * worker_span; + int const worker_end = imin(route_count, worker_begin + worker_span); + int const warp_span = ((worker_span + warps * 32 - 1) / (warps * 32)) * 32; + int const begin = worker_begin + warp * warp_span; + int const end = imin(worker_end, begin + warp_span); + int hot_count = shared[warp * bins + experts]; + hot_count = imin(hot_count, imax(0, end - begin)); + for (int offset = lane; offset < hot_count; offset += 32) + { + int index = hot_indices[begin + offset]; + int expert = routes[index]; + int b = expert_lut[expert]; + int ordinal = output[index] + shared[warp * bins + expert]; + int const* prefix = route_prefix + b * (ep + 1); + int destination = route_destination(prefix, ordinal, ep); + int slot = -1; + if (group_contains(destination, p_begin[b], p_size[b])) + { + slot = destination * local_slots + home + p_helper[b]; + } + else if (destination == p_owner[b]) + { + slot = destination * local_slots + expert - p_owner[b] * home; + } + else + { + record_status(status + 5, epoch); + } + output[index] = slot; + } +} + +__device__ void materialize_worker_routes(int const* routes, int* output, int const* route_prefix, int const* plan, + int* status, int route_count, int ep, int experts, int helpers, int worker, int workers, Epoch epoch) +{ + extern __shared__ int shared[]; + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const warps = blockDim.x >> 5; + int const bins = experts + 1; + int* expert_lut = worker_expert_lut(experts); + int const home = experts / ep; + int const local_slots = home + helpers; + int const* p_begin = plan + 2 + kPlanGroupBegin * kMaxBroadcasts; + int const* p_size = plan + 2 + kPlanGroupSize * kMaxBroadcasts; + int const* p_owner = plan + 2 + kPlanOwner * kMaxBroadcasts; + int const* p_helper = plan + 2 + kPlanHelper * kMaxBroadcasts; + + int const worker_span = (route_count + workers - 1) / workers; + int const worker_begin = worker * worker_span; + int const worker_end = imin(route_count, worker_begin + worker_span); + int const warp_span = ((worker_span + warps * 32 - 1) / (warps * 32)) * 32; + int const begin = worker_begin + warp * warp_span; + int const end = imin(worker_end, begin + warp_span); + for (int index = begin + lane; index < end; index += 32) + { + int expert = routes[index]; + int slot = -1; + if (expert >= 0 && expert < experts) + { + int entry = expert_lut[expert]; + if (entry < 0) + { + slot = -entry - 1; + } + else + { + int b = entry; + int ordinal = output[index] + shared[warp * bins + expert]; + int const* prefix = route_prefix + b * (ep + 1); + int destination = route_destination(prefix, ordinal, ep); + if (group_contains(destination, p_begin[b], p_size[b])) + { + slot = destination * local_slots + home + p_helper[b]; + } + else if (destination == p_owner[b]) + { + slot = destination * local_slots + expert - p_owner[b] * home; + } + else + { + record_status(status + 5, epoch); + } + } + } + else if (expert != -1) + { + record_status(status + 3, epoch); + } + output[index] = slot; + } +} + +template +__global__ __launch_bounds__(512, 1) void halo_q_scheduler_kernel(int const* routes, int* out_slots, int* out_ids, + int* out_levels, int* out_owners, unsigned long long const* peer_bases, int* status, int* partial, int* route_aux, + int* grid_sync, int* plan, int* route_prefix, int* plan_channel, int runtime_ep, int runtime_experts, + int runtime_helpers, int local_rank, int route_capacity, int ctas, int enable_pdl, unsigned long long spin_cycles, + int plan_abi_version, int plan_channel_words, int route_features, int capacity, int route_count) +{ + int const ep = FixedEp > 0 ? FixedEp : runtime_ep; + int const experts = FixedExperts > 0 ? FixedExperts : runtime_experts; + int const helpers = FixedHelpers > 0 ? FixedHelpers : runtime_helpers; + (void) plan_channel_words; // Validated by the host launch wrapper. + /* The scheduler grid is deliberately residency-safe (one CTA per SM and + * CTA count <= SM count), so signal as soon as every CTA starts. The + * independent same-stream quantizer may then occupy all unused SMs while + * these CTAs retain their own SMs for the complete scheduler lifetime. */ + if (enable_pdl && threadIdx.x == 0) + cudaTriggerProgrammaticLaunchCompletion(); + // Only the caller's valid prefix is readable. Keep the public output at + // its configured capacity, clearing stale routes inside this same kernel. + // Empty ranks still execute every collective; no route load occurs at N=0. + for (int index = route_count + blockIdx.x * blockDim.x + threadIdx.x; index < route_capacity; + index += gridDim.x * blockDim.x) + out_slots[index] = -1; + local_histogram_pass(routes, partial, route_count, experts, ctas); + Epoch epoch; + if (!grid_rendezvous(grid_sync, ctas, spin_cycles, status, &epoch)) + return; + + if (blockIdx.x == 0) + { + extern __shared__ int shared[]; + int offset = ep * experts; + int* counts = shared; + int* histogram = shared + offset; + offset += experts; + int* load = shared + offset; + offset += ep; + // Reuse the exact capacity already validated/computed by launch(). + int* quota = shared + offset; + offset += capacity * ep; + int* selected_flags = shared + offset; + offset += experts; + int* recv_used = shared + offset; + offset += ep * kHelperMaskWords; + int* send_used = shared + offset; + offset += ep; + int* next_index = shared + offset; + offset += ep; + int* frozen = shared + offset; + offset += ep; + int* ring = shared + offset; + offset += ep; + int* order = shared + offset; + offset += capacity; + int* host_mask = shared + offset; + + bool exchange_ok; + if (ep == 8) + exchange_ok = exchange_counts<8>( + peer_bases, partial, shared, ep, experts, local_rank, ctas, epoch, spin_cycles, status); + else + exchange_ok = exchange_counts<0>( + peer_bases, partial, shared, ep, experts, local_rank, ctas, epoch, spin_cycles, status); + if (!exchange_ok) + return; + /* EP8 keeps only its candidate masses; other EPs retain the dense + * histogram required by the generic planner path. */ + if (ep == 8 && helpers == 4 && experts >= 32) + { + prepare_halo_m_sparse8(counts, histogram, load, selected_flags, experts); + } + else + { + prepare_halo_m<0, 0>(counts, histogram, load, selected_flags, recv_used, send_used, next_index, frozen, ep, + experts, helpers); + } + + int broadcasts = 0; + if (threadIdx.x < (ep <= 8 ? 8 : 32)) + { + // These short-circuit bounds keep helpers*ep <= 120. + if (ep <= 8 && helpers <= 15 && helpers * ep <= experts) + { + if (ep == 8 && helpers == 4) + { + broadcasts = plan_halo_m_register8<8, 4>(histogram, load, quota, selected_flags, plan, ep, helpers); + } + else if (ep == 4 && helpers == 4) + { + broadcasts = plan_halo_m_register8<4, 4>(histogram, load, quota, selected_flags, plan, ep, helpers); + } + else + { + broadcasts = plan_halo_m_register8<0, 0>(histogram, load, quota, selected_flags, plan, ep, helpers); + } + } + else + { + broadcasts = plan_halo_m_warp(histogram, load, quota, selected_flags, recv_used, send_used, next_index, + frozen, plan, ep, experts, helpers); + } + if (threadIdx.x == 0) + plan[1] = broadcasts; + } + __syncthreads(); + broadcasts = plan[1]; + /* p_expert and plan[1] are produced by thread zero. Publish this + * placement-only epoch before HALO-Q so workers can construct their + * selected-expert LUT while quota repair is still running. */ + if (ctas > 1 && threadIdx.x == 0) + { + store_release_device(plan + 0, epoch); + } + /* Small-EP repair excludes warp8 and can overlap coloring. Generic + * repair first meets the coloring warp at its full-block barrier. */ + if (threadIdx.x == 8 * 32) + color_and_publish_local( + plan, broadcasts, ep, helpers, local_rank, recv_used, out_ids, out_levels, out_owners, status, epoch); + if (ep == 8) + run_halo_q<8>(counts, quota, plan, broadcasts, ep, experts, ring, order, host_mask); + else + run_halo_q<0>(counts, quota, plan, broadcasts, ep, experts, ring, order, host_mask); + if (ctas > 1) + { + int const warp = threadIdx.x >> 5; + /* build_route_prefix is independent of warp 8's outputs. Defer their + * rendezvous until both paths reach their first shared consumer so + * prefix construction and coloring can overlap safely. */ + if (warp == 0) + { + build_route_prefix(counts, quota, route_prefix, plan, broadcasts, ep, experts, local_rank, 32); + __syncwarp(); + } + if (warp == 0 || warp == 8) + scheduler_tail_barrier(); + if (warp == 0) + { + /* Release the worker CTAs here rather than after the channel + * publish. They gate on grid_sync+2 and consume route_prefix, + * which is complete above; plan_channel is consumed only by the + * copy stack, which acquires on the channel's own epoch word + * written with store_release_system below. Publishing after + * the release lets the workers coloring pass -- the work that + * actually determines when the kernel ends -- overlap the whole + * publish instead of waiting behind it. */ + if (threadIdx.x == 0) + store_release_device(grid_sync + 2, epoch); + sami_copy::publish_plan_channel_warp(plan_channel, out_ids, out_levels, out_owners, + copy_placement_view(plan), helpers, ep, local_rank, experts, plan_abi_version, route_features); + __syncwarp(); + } + return; + } + if (threadIdx.x == 0) + store_release_device(plan + 0, epoch); + __syncthreads(); + build_route_prefix(counts, quota, route_prefix, plan, broadcasts, ep, experts, local_rank, blockDim.x); + __syncthreads(); + sami_copy::publish_plan_channel_warp(plan_channel, out_ids, out_levels, out_owners, copy_placement_view(plan), + helpers, ep, local_rank, experts, plan_abi_version, route_features); + /* The block barrier orders route-prefix/helper writes before release. */ + __syncthreads(); + } + else + { + int const worker = blockIdx.x - 1; + int const workers = ctas - 1; + if (!stable_worker_ordinal_pass( + routes, out_slots, route_aux, status, route_count, experts, worker, workers, spin_cycles)) + return; + int* expert_lut = worker_expert_lut(experts); + prepare_worker_default_slots(expert_lut, ep, experts, helpers); + __shared__ int plan_ready; + if (!wait_for_plan_epoch(plan + 0, epoch, spin_cycles, status, &plan_ready)) + return; + prepare_worker_selected_experts(expert_lut, plan); + pre_materialize_and_compact_worker_routes( + routes, out_slots, route_aux, status, route_count, experts, worker, workers, epoch); + if (!wait_for_plan_epoch(grid_sync + 2, epoch, spin_cycles, status, &plan_ready)) + return; + materialize_worker_hot_routes(routes, out_slots, route_prefix, plan, route_aux, status, route_count, ep, + experts, helpers, worker, workers, epoch); + return; + } + + /* CTA-count one remains a correct serial fallback. */ + if (!stable_worker_ordinal_pass(routes, out_slots, route_aux, status, route_count, experts, 0, 1, spin_cycles)) + return; + int* expert_lut = worker_expert_lut(experts); + prepare_worker_default_slots(expert_lut, ep, experts, helpers); + prepare_worker_selected_experts(expert_lut, plan); + materialize_worker_routes( + routes, out_slots, route_prefix, plan, status, route_count, ep, experts, helpers, 0, 1, epoch); +} + +bool validPlanChannelLaunch(std::uintptr_t pointer, int abiVersion, int words, int features, int ep, int helpers) +{ + if (pointer == 0) + { + return abiVersion == 0 && words == 0 && features == 0; + } + if (pointer % 16 != 0 || helpers <= 0) + { + return false; + } + auto const stride = (static_cast(helpers) + 3) & ~std::int64_t{3}; + auto const expected = 4 + 7 * stride; + return abiVersion == 6 && ep >= 2 && ep <= kMaxEp && ep % 2 == 0 && (features == 0 || features == 1) + && expected <= INT_MAX && words == expected; +} + +cudaError_t ensureDynamicSharedMemory(void const* kernel, int sharedBytes) +{ + int device = -1; + auto error = cudaGetDevice(&device); + if (error != cudaSuccess) + return error; + + using KernelCache = std::map>; + static std::mutex mutex; + static std::map configuredBytes; + std::lock_guard lock(mutex); + auto& current = configuredBytes[device][kernel]; + if (current >= sharedBytes) + return cudaSuccess; + error = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, sharedBytes); + if (error == cudaSuccess) + current = sharedBytes; + return error; +} + +} // namespace + +void invokeMoeRebalanceHaloQ(MoeRebalanceHaloQParams const& params, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO((params.routes != nullptr || params.validRouteCount == 0) && params.outSlots != nullptr + && params.outIds != nullptr && params.outLevels != nullptr && params.outOwners != nullptr + && params.peerBases != nullptr && params.status != nullptr && params.partial != nullptr + && params.routeAux != nullptr && params.gridSync != nullptr && params.planWorkspace != nullptr + && params.routePrefix != nullptr && params.ep >= 2 && params.ep <= kMaxEp && params.ep % 2 == 0 + && params.experts > 0 && params.experts <= kMaxExperts && params.experts % params.ep == 0 + && params.helpers > 0 && params.helpers <= INT_MAX / params.ep - params.experts / params.ep + && params.localRank >= 0 && params.localRank < params.ep && params.validRouteCount >= 0 + && params.validRouteCount <= params.routeCapacity && params.routeCapacity > 0 + && params.routeCapacity <= INT_MAX / params.ep && params.routeCapacity <= INT_MAX - params.ctas + && params.ctas > 0 && params.ctas <= 128 && params.threads == 512, + "Invalid HALO-Q scheduler launch parameters"); + TLLM_CHECK_WITH_INFO( + validPlanChannelLaunch(reinterpret_cast(params.planChannel), params.planAbiVersion, + params.planChannelWords, params.routeFeatures, params.ep, params.helpers), + "Invalid HALO-Q PlanChannel launch ABI"); + + int const warps = params.threads / 32; + int const passInts = warps * (params.experts + 1); + int const workerInts = passInts + params.experts; + int capacity = broadcast_capacity(params.ep, params.experts, params.helpers); + int const plannerInts = params.ep * params.experts + params.experts + params.ep + capacity * params.ep + + params.experts + (kHelperMaskWords + 4) * params.ep + 2 * capacity; + std::size_t const sharedBytes + = static_cast(workerInts > plannerInts ? workerInts : plannerInts) * sizeof(int); + + void const* kernel = reinterpret_cast(&halo_q_scheduler_kernel<0, 0, 0>); + if (params.experts == 384 && params.helpers == 4) + { + if (params.ep == 8) + { + kernel = reinterpret_cast(&halo_q_scheduler_kernel<8, 384, 4>); + } + else if (params.ep == 4) + { + kernel = reinterpret_cast(&halo_q_scheduler_kernel<4, 384, 4>); + } + } + if (sharedBytes > 48 * 1024) + { + TLLM_CUDA_CHECK(ensureDynamicSharedMemory(kernel, static_cast(sharedBytes))); + } + + auto routes = params.routes; + auto outSlots = params.outSlots; + auto outIds = params.outIds; + auto outLevels = params.outLevels; + auto outOwners = params.outOwners; + auto peerBases = params.peerBases; + auto status = params.status; + auto partial = params.partial; + auto routeAux = params.routeAux; + auto gridSync = params.gridSync; + auto planWorkspace = params.planWorkspace; + auto routePrefix = params.routePrefix; + auto planChannel = params.planChannel; + int ep = params.ep; + int experts = params.experts; + int helpers = params.helpers; + int localRank = params.localRank; + int routeCapacity = params.routeCapacity; + int ctas = params.ctas; + int enablePdl = params.enablePdl ? 1 : 0; + std::uint64_t spinCycles = params.spinCycles; + int planAbiVersion = params.planAbiVersion; + int planChannelWords = params.planChannelWords; + int routeFeatures = params.routeFeatures; + int validRouteCount = params.validRouteCount; + void* arguments[] + = {&routes, &outSlots, &outIds, &outLevels, &outOwners, &peerBases, &status, &partial, &routeAux, &gridSync, + &planWorkspace, &routePrefix, &planChannel, &ep, &experts, &helpers, &localRank, &routeCapacity, &ctas, + &enablePdl, &spinCycles, &planAbiVersion, &planChannelWords, &routeFeatures, &capacity, &validRouteCount}; + TLLM_CUDA_CHECK( + cudaLaunchKernel(kernel, dim3(ctas, 1, 1), dim3(params.threads, 1, 1), arguments, sharedBytes, stream)); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h new file mode 100644 index 000000000000..d0215ab74617 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h @@ -0,0 +1,67 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef TRTLLM_KERNELS_MOE_REBALANCE_HALO_Q_H +#define TRTLLM_KERNELS_MOE_REBALANCE_HALO_Q_H + +#include "tensorrt_llm/common/config.h" + +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +struct MoeRebalanceHaloQParams +{ + int const* routes; + int* outSlots; + int* outIds; + int* outLevels; + int* outOwners; + std::uint64_t const* peerBases; + int* status; + int* partial; + int* routeAux; + int* gridSync; + int* planWorkspace; + int* routePrefix; + int* planChannel; + int ep; + int experts; + int helpers; + int localRank; + int routeCapacity; + int ctas; + int threads; + bool enablePdl; + std::uint64_t spinCycles; + int planAbiVersion; + int planChannelWords; + int routeFeatures; + int validRouteCount; +}; + +void invokeMoeRebalanceHaloQ(MoeRebalanceHaloQParams const& params, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END + +#endif // TRTLLM_KERNELS_MOE_REBALANCE_HALO_Q_H diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlaceholder.cpp b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlaceholder.cpp new file mode 100644 index 000000000000..e5afc3c3e2d4 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlaceholder.cpp @@ -0,0 +1,100 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h" +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h" + +#include "tensorrt_llm/common/assert.h" + +#include + +#ifdef TRTLLM_DYNAMIC_EPLB_PLACEHOLDER_HALO_Q +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +void invokeMoeRebalanceHaloQ(MoeRebalanceHaloQParams const&, cudaStream_t) +{ + TLLM_THROW("Dynamic EPLB HALO-Q requires a requested SM100, SM103, or SM107 target"); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END +#endif + +#ifdef TRTLLM_DYNAMIC_EPLB_PLACEHOLDER_TMA +namespace +{ + +int unsupported() +{ + return static_cast(cudaErrorNotSupported); +} + +} // namespace + +extern "C" int megamoe_tma_copy_create(uint64_t, int, int, MegamoeTmaCopyState** out) +{ + if (out != nullptr) + { + *out = nullptr; + } + return unsupported(); +} + +extern "C" int megamoe_tma_copy_configure_gpu_plan(MegamoeTmaCopyState*, MegamoeTmaGpuPlanConfig const*) +{ + return unsupported(); +} + +extern "C" int megamoe_tma_copy_bind_gpu_direct( + MegamoeTmaCopyState*, uint64_t, uint64_t, uint64_t, uint64_t, uint64_t, uint64_t, int) +{ + return unsupported(); +} + +extern "C" int megamoe_tma_copy_submit_gpu_direct(MegamoeTmaCopyState*, uint64_t, uint64_t, void*) +{ + return unsupported(); +} + +extern "C" int megamoe_tma_copy_gpu_plan_result(MegamoeTmaCopyState*, MegamoeTmaGpuPlanResult*) +{ + return unsupported(); +} + +extern "C" int megamoe_tma_copy_destroy(MegamoeTmaCopyState** state) +{ + return state == nullptr || *state == nullptr ? static_cast(cudaSuccess) : unsupported(); +} + +extern "C" int megamoe_tma_copy_config_info(MegamoeTmaCopyState const*, MegamoeTmaCopyConfig*) +{ + return unsupported(); +} + +extern "C" char const* megamoe_tma_copy_error_string(int error) +{ + if (error == static_cast(cudaErrorNotSupported)) + { + return "Dynamic EPLB TMA requires CUDA 13.1+ and a requested SM100, SM103, or SM107 target"; + } + return cudaGetErrorString(static_cast(error)); +} +#endif diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlanPolicy.h b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlanPolicy.h new file mode 100644 index 000000000000..b577ad28d2cc --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlanPolicy.h @@ -0,0 +1,380 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef TRTLLM_KERNELS_MOE_REBALANCE_PLAN_POLICY_H +#define TRTLLM_KERNELS_MOE_REBALANCE_PLAN_POLICY_H + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +// Copy-owned GPU policy adapter. The scheduler supplies only placement +// geometry. +// Fusing this adapter into its publication tail avoids another kernel launch or +// any CPU policy work. Without a bound copy channel, the adapter does nothing. +namespace sami_copy +{ +constexpr int kMaxEp = 32; +constexpr int kMinGroup = 4; + +struct PlacementView +{ + int count, capacity; + int const *experts, *levels, *begins, *sizes, *owners, *helpers; +}; + +__device__ __forceinline__ int imin(int a, int b) +{ + return a < b ? a : b; +} + +__device__ __forceinline__ int imax(int a, int b) +{ + return a > b ? a : b; +} + +__device__ __forceinline__ bool group_contains(int rank, int begin, int size) +{ + return rank >= begin && rank < begin + size; +} + +__host__ __device__ __forceinline__ int hierarchy_depth(int ep) +{ + int depth = 1; + for (int level = 1; (1 << level) <= ep; ++level) + { + int groups = 1 << level; + if (ep % groups || ep / groups < kMinGroup) + break; + ++depth; + } + return depth; +} + +__device__ __forceinline__ void store_release_system(int* address, int value) +{ + cuda::atomic_ref atom(*address); + atom.store(value, cuda::memory_order_release); +} + +/* ABI6 publishes explicit per-slot modes and bounded owner jobs. All scans + * cover the complete global plan, including broadcasts beyond one warp. Only + * the releasing lane stores mapped words, preserving the SYS release contract. + */ +template +__device__ void publish_plan_channel(int* channel, int const* ids, int const* levels, int const* owners, + PlacementView const& plan, int helpers, int ep, int local_rank, int expert_count, int plan_abi_version, + int route_features) +{ + if (!channel || threadIdx.x != 0) + return; + int const stride = (helpers + 3) & ~3; + int const broadcasts = plan.count; + int const* p_expert = plan.experts; + int const* p_level = plan.levels; + int const* p_begin = plan.begins; + int const* p_size = plan.sizes; + int const* p_owner = plan.owners; + int const* p_helper = plan.helpers; + int owner_count[kMaxEp] = {0}; + int status = plan_abi_version != 6 || (route_features & ~1) || broadcasts < 0 || broadcasts > plan.capacity; + int global_max_send = 0; + int common_size = broadcasts > 0 ? p_size[0] : 0; + int foreign_begin = -1, target_active = 0; + bool common_proper_size = common_size > 0 && common_size < ep; + bool one_foreign_target = true; + int const home = expert_count / ep; + for (int b = 0; b < broadcasts && !status; ++b) + { + int const level = p_level[b]; + int const size = level >= 0 && level < hierarchy_depth(ep) ? ep / (1 << level) : 0; + if (p_expert[b] < 0 || p_expert[b] >= expert_count || p_owner[b] != p_expert[b] / home || !size + || p_size[b] != size || p_begin[b] < 0 || p_begin[b] > ep - size || p_begin[b] % size || p_helper[b] < 0 + || p_helper[b] >= helpers) + { + status = 1; + break; + } + for (int j = 0; j < b; ++j) + { + bool const intersects = p_begin[j] < p_begin[b] + size && p_begin[b] < p_begin[j] + p_size[j]; + if (p_expert[j] == p_expert[b] || (intersects && p_helper[j] == p_helper[b])) + status = 1; + } + if (group_contains(local_rank, p_begin[b], size) + && (ids[p_helper[b]] != p_expert[b] || levels[p_helper[b]] != level || owners[p_helper[b]] != p_owner[b])) + status = 1; + int const count = ++owner_count[p_owner[b]]; + if (count > imin(helpers, home)) + status = 1; + global_max_send = imax(global_max_send, count); + common_proper_size &= size == common_size; + if (!group_contains(p_owner[b], p_begin[b], size)) + { + if (foreign_begin < 0) + foreign_begin = p_begin[b]; + else + one_foreign_target &= foreign_begin == p_begin[b]; + } + } + if (common_proper_size && foreign_begin >= 0) + for (int b = 0; b < broadcasts; ++b) + target_active += p_begin[b] == foreign_begin && p_size[b] == common_size; + bool const external = !status && (route_features & 1) && common_proper_size && one_foreign_target + && foreign_begin >= 0 && (global_max_send <= 1 || (global_max_send <= 2 && target_active >= 3)); + + for (int slot = 0; slot < stride; ++slot) + { + bool const inactive = status || slot >= helpers; + channel[4 + slot] = inactive ? -1 : ids[slot]; + channel[4 + stride + slot] = inactive ? -1 : levels[slot]; + channel[4 + 2 * stride + slot] = inactive ? -1 : owners[slot]; + channel[4 + 3 * stride + slot] = 0; + channel[4 + 4 * stride + slot] = -1; + channel[4 + 5 * stride + slot] = -1; + channel[4 + 6 * stride + slot] = -1; + } + int outgoing = 0; + for (int b = 0; b < broadcasts && !status; ++b) + { + int const owner = p_owner[b], begin = p_begin[b], size = p_size[b]; + bool const local_owner = group_contains(owner, begin, size); + int mode = 0; + if (local_owner) + { + bool pure = true, best = true; + for (int j = 0; j < broadcasts; ++j) + { + bool const intersects = p_begin[j] < begin + size && begin < p_begin[j] + p_size[j]; + if (intersects + && (p_begin[j] != begin || p_size[j] != size || !group_contains(p_owner[j], begin, size))) + pure = false; + if (p_owner[j] == owner && group_contains(owner, p_begin[j], p_size[j]) + && (p_size[j] > size || (p_size[j] == size && j < b))) + best = false; + } + if (pure || best) + mode = 1; + } + else if (external) + { + mode = 2; + if (owner == local_rank) + { + int target_id = begin / size; + for (int level = 1; level < p_level[b]; ++level) + target_id += 1 << level; + channel[4 + 4 * stride + outgoing] = p_expert[b]; + channel[4 + 5 * stride + outgoing] = target_id; + channel[4 + 6 * stride + outgoing] = p_helper[b]; + ++outgoing; + } + } + if (group_contains(local_rank, begin, size)) + channel[4 + 3 * stride + p_helper[b]] = mode; + } + channel[1] = global_max_send; + channel[2] = status; + channel[3] = outgoing; + if constexpr (PublishHost) + store_release_system(channel, 0); + else + channel[0] = 0; +} + +/* Small ABI6 plans fit in one warp. Load each record once and compare its + * register values through shuffles. Two mode bits per helper fit in a register + * for S<=16. Larger/malformed plans use the unchanged scalar publisher. + * Only lane 0 writes mapped words; its SYS release orders every field. */ +template +__device__ void publish_plan_channel_warp(int* channel, int const* ids, int const* levels, int const* owners, + PlacementView const& plan, int helpers, int ep, int local_rank, int expert_count, int plan_abi_version, + int route_features) +{ + if (!channel || (threadIdx.x >> 5) != 0) + return; + // One shared cold fallback keeps its scalar O(B^2) implementation out + // of the hot sections instead of inlining a copy at every failure gate. + do + { + int const broadcasts = plan.count; + if (broadcasts < 0 || broadcasts > plan.capacity || broadcasts > 32 || helpers <= 0 || helpers > 16 + || plan_abi_version != 6 || (route_features & ~1)) + { + break; + } + constexpr unsigned mask = 0xffffffffu; + int const lane = threadIdx.x & 31; + bool const active = lane < broadcasts; + int const expert = active ? plan.experts[lane] : -1; + int const level = active ? plan.levels[lane] : 0; + int const begin = active ? plan.begins[lane] : 0; + int const size = active ? plan.sizes[lane] : 0; + int const owner = active ? plan.owners[lane] : -1; + int const helper = active ? plan.helpers[lane] : -1; + // Coalesced helper reads avoid three dependent, short-circuited global + // loads for every local record. Invalid helpers select a harmless lane + // and are rejected below before any publication. + int const slot_id = lane < helpers ? ids[lane] : -1; + int const slot_level = lane < helpers ? levels[lane] : -1; + int const slot_owner = lane < helpers ? owners[lane] : -1; + int const check_id = __shfl_sync(mask, slot_id, helper & 31); + int const check_level = __shfl_sync(mask, slot_level, helper & 31); + int const check_owner = __shfl_sync(mask, slot_owner, helper & 31); + int const home = expert_count / ep; + bool const level_ok = static_cast(level) < static_cast(hierarchy_depth(ep)); + int const expected_size = level_ok ? ep >> level : 0; + bool const aligned + = (ep & (ep - 1)) == 0 ? (begin & (expected_size - 1)) == 0 : begin % imax(expected_size, 1) == 0; + // Unsigned interval checks stay defined even for malformed int32 fields. + // The native launch contract guarantees positive home and even EP<=32. + bool const owner_ok = (static_cast(owner) < static_cast(ep)) + & (static_cast(expert) - static_cast(owner) * static_cast(home) + < static_cast(home)); + bool const local + = static_cast(local_rank) - static_cast(begin) < static_cast(expected_size); + bool const record_valid = (expert >= 0) & (expert < expert_count) & owner_ok & (expected_size > 0) + & (size == expected_size) & (begin >= 0) & (begin <= ep - expected_size) & aligned & (helper >= 0) + & (helper < helpers); + bool const local_valid = (check_id == expert) & (check_level == level) & (check_owner == owner); + bool const valid = !active | (record_valid & (!local | local_valid)); + // Reject invalid addresses before the pairwise geometry calculations. + if (__any_sync(mask, !valid)) + { + break; + } + unsigned const same_owner = __match_any_sync(mask, owner); + int const owned = active ? __popc(same_owner) : 0; + int const global_max_send = __reduce_max_sync(mask, owned); + unsigned const same_expert = __match_any_sync(mask, expert); + int bad = active & ((__popc(same_expert) != 1) | (owned > imin(helpers, home))); + int pure = 1, best = 1; + // Validated EP<=32/S<=16 geometry needs only 19 bits. Exchange one + // packed register per pair rather than four independent shuffles. + unsigned const pair_record = static_cast(begin) | (static_cast(size - 1) << 5) + | (static_cast(owner) << 10) | (static_cast(helper) << 15); +// Eager bitwise predicates avoid divergent short-circuit branches and +// repeated reconvergence inside this register-only comparison loop. +// All lanes execute every shuffle, including inactive plan lanes. +#pragma unroll 1 + for (int j = 0; j < broadcasts; ++j) + { + unsigned const other = __shfl_sync(mask, pair_record, j); + int const other_begin = other & 31u; + int const other_size = ((other >> 5) & 31u) + 1; + int const other_owner = (other >> 10) & 31u; + int const other_helper = (other >> 15) & 15u; + int const intersects = (other_begin < begin + size) & (begin < other_begin + other_size); + bad |= active & (j < lane) & intersects & (other_helper == helper); + pure &= !intersects + | ((other_begin == begin) & (other_size == size) & (other_owner >= begin) + & (other_owner < begin + size)); + best &= !((other_owner == owner) & (owner >= other_begin) & (owner < other_begin + other_size) + & ((other_size > size) | ((other_size == size) & (j < lane)))); + } + if (__any_sync(mask, bad)) + { + // Preserve even the scalar error header/partial owner-count semantics. + break; + } + int const common_size = __shfl_sync(mask, size, 0); + bool const common_proper_size + = broadcasts > 0 && common_size < ep && !__any_sync(mask, active && size != common_size); + bool const local_owner = active && group_contains(owner, begin, size); + bool const foreign = active && !local_owner; + unsigned const foreign_mask = __ballot_sync(mask, foreign); + int const first_foreign = foreign_mask ? __ffs(foreign_mask) - 1 : 0; + int const foreign_begin = __shfl_sync(mask, begin, first_foreign); + bool const one_foreign_target = !__any_sync(mask, foreign && begin != foreign_begin); + int const target_active = __popc(__ballot_sync(mask, active && begin == foreign_begin && size == common_size)); + bool const external = (route_features & 1) && common_proper_size && one_foreign_target && foreign_mask + && (global_max_send <= 1 || (global_max_send <= 2 && target_active >= 3)); + int const mode = local_owner ? (pure | best) : (foreign && external ? 2 : 0); + unsigned const local_mode + = active && group_contains(local_rank, begin, size) ? static_cast(mode) << (2 * helper) : 0; + unsigned const mode_bits = __reduce_or_sync(mask, local_mode); + unsigned const outgoing_mask = __ballot_sync(mask, mode == 2 && owner == local_rank); + int const stride = (helpers + 3) & ~3; + // Reuse the coalesced helper registers from validation. All lanes + // participate in every shuffle; lane 0 remains the only mapped writer. + for (int slot = 0; slot < stride; slot += 4) + { + const int4 ids4 = make_int4(__shfl_sync(mask, slot_id, slot), __shfl_sync(mask, slot_id, slot + 1), + __shfl_sync(mask, slot_id, slot + 2), __shfl_sync(mask, slot_id, slot + 3)); + const int4 levels4 = make_int4(__shfl_sync(mask, slot_level, slot), __shfl_sync(mask, slot_level, slot + 1), + __shfl_sync(mask, slot_level, slot + 2), __shfl_sync(mask, slot_level, slot + 3)); + const int4 owners4 = make_int4(__shfl_sync(mask, slot_owner, slot), __shfl_sync(mask, slot_owner, slot + 1), + __shfl_sync(mask, slot_owner, slot + 2), __shfl_sync(mask, slot_owner, slot + 3)); + if (lane == 0) + { + // Native channel/stride alignment permits 16-byte stores. + const int4 empty = make_int4(-1, -1, -1, -1); + *reinterpret_cast(channel + 4 + slot) = ids4; + *reinterpret_cast(channel + 4 + stride + slot) = levels4; + *reinterpret_cast(channel + 4 + 2 * stride + slot) = owners4; + unsigned const packed = mode_bits >> (2 * slot); + *reinterpret_cast(channel + 4 + 3 * stride + slot) + = make_int4(packed & 3u, (packed >> 2) & 3u, (packed >> 4) & 3u, (packed >> 6) & 3u); + *reinterpret_cast(channel + 4 + 4 * stride + slot) = empty; + *reinterpret_cast(channel + 4 + 5 * stride + slot) = empty; + *reinterpret_cast(channel + 4 + 6 * stride + slot) = empty; + } + } + // Canonical target IDs enumerate proper hierarchy levels, widest first. + int const target = mode == 2 ? begin / size + (1 << level) - 2 : -1; + int outgoing = 0; + for (unsigned remaining = outgoing_mask; remaining; remaining &= remaining - 1) + { + int const source_lane = __ffs(remaining) - 1; + int const out_expert = __shfl_sync(mask, expert, source_lane); + int const out_target = __shfl_sync(mask, target, source_lane); + int const out_helper = __shfl_sync(mask, helper, source_lane); + if (lane == 0) + { + channel[4 + 4 * stride + outgoing] = out_expert; + channel[4 + 5 * stride + outgoing] = out_target; + channel[4 + 6 * stride + outgoing] = out_helper; + } + ++outgoing; + } + if (lane == 0) + { + channel[1] = global_max_send; + channel[2] = 0; + channel[3] = outgoing; + if constexpr (PublishHost) + store_release_system(channel, 0); + else + channel[0] = 0; + } + return; + } while (false); + publish_plan_channel( + channel, ids, levels, owners, plan, helpers, ep, local_rank, expert_count, plan_abi_version, route_features); +} + +} // namespace sami_copy + +} // namespace kernels + +TRTLLM_NAMESPACE_END + +#endif // TRTLLM_KERNELS_MOE_REBALANCE_PLAN_POLICY_H diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu new file mode 100644 index 000000000000..bef362a9bdc6 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu @@ -0,0 +1,1085 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +/* SM-driven unicast -> multicast copy, PTX 9.1 / CUDA 13.1 or newer. + * One aligned TMA load handles each eligible 8 KiB logical slice. Physical + * scratch slots include 16 bytes of padding for aligned head/tail accesses. + * + * Each communication warp owns a near-equal share of all whole 8 KiB slots, split into two + * banks. A bank is one bulk group containing all of its 8 KiB slice slots. + * A producer warp issues UC loads; its consumer warp issues MC stores. + * Per-bank alternating tickets protect slot metadata and shared-memory reuse. + * This permits every allocated slot to participate even when one warp owns + * many slots, without exceeding the ISA's outstanding bulk-group limit. + */ +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlanPolicy.h" +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h" +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlanBuilder.cuh" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#if CUDART_VERSION < 13010 +#error "TMA multicast copy requires CUDA >= 13.1 (PTX >= 9.1)" +#endif +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 900 +#error "TMA multicast copy requires SM90 or newer" +#endif + +namespace +{ +constexpr unsigned kSlice = MEGAMOE_TMA_COPY_SLICE_BYTES; +constexpr unsigned kSlotControlBytes = MEGAMOE_TMA_COPY_SLOT_CONTROL_BYTES; +constexpr unsigned kSlotDataBytes = MEGAMOE_TMA_COPY_SLOT_DATA_BYTES; +constexpr unsigned kCoalesced = 4U; +constexpr unsigned kFallbackBulk = 8U; + +struct alignas(16) SlotControl +{ + unsigned long long barrier; + unsigned phase; + unsigned reserved; // head=1, tail=2, coalesced=4, fallback-has-bulk=8 + uint64_t fast_dst; + unsigned fast_bytes; + unsigned fast_padding; // first slot in each bank: even=free, odd=ready ticket +}; + +static_assert(kSlotControlBytes == 32, "cached consume requires 32-byte control"); +static_assert(offsetof(SlotControl, fast_dst) == 16, "cached destination alignment"); +static_assert(offsetof(SlotControl, fast_bytes) == 24, "cached byte-count offset"); +static_assert(kSlotDataBytes == kSlice + 16, "one aligned edge of slot padding"); +static_assert(sizeof(SlotControl) == kSlotControlBytes, "shared slot control ABI"); +static_assert(kSlotDataBytes % 16 == 0, "TMA slot alignment"); +static_assert(sizeof(MegamoeTmaCopySegment) == 4 * sizeof(uint64_t), "segment ABI"); +static_assert(sizeof(MegamoeTmaCopyRange) == 9 * sizeof(uint64_t), "range ABI"); +static_assert(sizeof(MegamoeTmaCopyConfig) == 88, "configuration ABI"); + +__device__ __forceinline__ unsigned shared_address(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ void initialize_barrier(SlotControl* slot) +{ + unsigned bar = shared_address(&slot->barrier); + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(bar) : "memory"); + slot->phase = 0; + slot->reserved = 0; + slot->fast_padding = 0; +} + +__device__ __forceinline__ void expect_slice(SlotControl* slot, unsigned bytes) +{ + unsigned bar = shared_address(&slot->barrier); + unsigned long long arrival; + asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 %0, [%1], %2;" + : "=l"(arrival) + : "r"(bar), "r"(bytes) + : "memory"); + asm volatile("" ::"l"(arrival)); +} + +__device__ __forceinline__ void load_fragment(unsigned char* buffer, SlotControl* slot, uint64_t src, unsigned bytes) +{ + unsigned bar = shared_address(&slot->barrier); + unsigned smem = shared_address(buffer); + asm volatile( + "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes " + "[%0], [%1], %2, [%3];" ::"r"(smem), + "l"(src), "r"(bytes), "r"(bar) + : "memory"); +} + +__device__ __forceinline__ void await_slice(SlotControl* slot) +{ + unsigned ready; + unsigned bar = shared_address(&slot->barrier); + unsigned phase = slot->phase; + do + { + asm volatile( + "{ .reg .pred p; " + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 p, [%1], %2; " + "selp.u32 %0, 1, 0, p; }" + : "=r"(ready) + : "r"(bar), "r"(phase) + : "memory"); + if (!ready) + __nanosleep(32); + } while (!ready); + // Only a slot that actually issued and completed TMA advances its phase. + slot->phase = phase ^ 1U; +} + +__device__ __forceinline__ void multicast_slice(uint64_t dst, unsigned char const* buffer, unsigned bytes) +{ + unsigned smem = shared_address(buffer); + asm volatile( + "multimem.cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;" ::"l"(dst), "r"(smem), "r"(bytes) + : "memory"); +} + +__device__ __forceinline__ void multicast_edge(uint64_t dst, unsigned char const* buffer, uint16_t mask) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + unsigned smem = shared_address(buffer); + asm volatile( + "multimem.cp.async.bulk.global.shared::cta.bulk_group.cp_mask " + "[%0], [%1], 16, %2;" ::"l"(dst), + "r"(smem), "h"(mask) + : "memory"); +#else + // coalesced_copy_flags() is zero on these targets; never use .cp_mask on SM90. + asm volatile("trap;" ::: "memory"); +#endif +} + +__device__ __forceinline__ void wait_older_reads() +{ + asm volatile("cp.async.bulk.wait_group.read 1;" ::: "memory"); +} + +__device__ __forceinline__ void wait_all_reads() +{ + // Buffer reuse only requires completion of the SMEM reads. Destination + // writes are drained separately before publishing READY. + asm volatile("cp.async.bulk.wait_group.read 0;" ::: "memory"); +} + +__device__ __forceinline__ void wait_all_writes() +{ + // Deliberately no .read: the multicast destination writes must be complete. + asm volatile("cp.async.bulk.wait_group 0;" ::: "memory"); +} + +__device__ __forceinline__ void system_publish_fence() +{ + asm volatile("fence.proxy.alias;" ::: "memory"); + asm volatile("fence.release.sys;" ::: "memory"); +} + +__device__ __forceinline__ void multicast_generation(uint64_t flag_mc, uint64_t generation) +{ + // The notification uses the same MC backing/UC acquire protocol as the + // former copy-engine terminal. Publish only after all issuers' writes. + system_publish_fence(); + asm volatile("multimem.st.release.sys.global.u64 [%0], %1;" ::"l"(flag_mc), "l"(generation) : "memory"); + asm volatile("fence.proxy.alias;" ::: "memory"); +} + +__device__ __forceinline__ void finish_cta_and_notify(unsigned* completed_ctas, uint64_t flag_mc, uint64_t generation) +{ + unsigned previous; + // A release/acquire RMW chain combines the completed payload of all CTAs. + // No CTA waits for another: oversubscribed grids cannot deadlock here. + asm volatile("atom.acq_rel.gpu.global.add.u32 %0, [%1], 1;" : "=r"(previous) : "l"(completed_ctas) : "memory"); + if (previous == gridDim.x - 1) + { + // All RMWs are done. The next launch is ordered after this kernel, so + // the last CTA can rearm the counter without a host memset each time. + asm volatile("st.relaxed.gpu.global.u32 [%0], 0;" ::"l"(completed_ctas) : "memory"); + multicast_generation(flag_mc, generation); + } +} + +struct Slice +{ + uint64_t begin; + uint64_t end; + uint64_t segment_begin; + uint64_t segment_end; +}; + +struct Fragment +{ + uint64_t src; + uint64_t dst; + unsigned bytes; + unsigned scalar_prefix; + unsigned bulk_bytes; + unsigned smem_offset; +}; + +__device__ __forceinline__ Slice locate_slice( + MegamoeTmaCopySegment const* segments, MegamoeTmaCopyRange const* ranges, uint64_t range_count, uint64_t ordinal) +{ + uint64_t lo = 0, hi = range_count; + while (lo < hi) + { + uint64_t mid = lo + (hi - lo) / 2; + if (ordinal < ranges[mid].prefix_end) + hi = mid; + else + lo = mid + 1; + } + const MegamoeTmaCopyRange range = ranges[lo]; + const uint64_t local = ordinal - (range.prefix_end - range.slice_count); + uint64_t slice = range.first_slice + local * range.slice_stride; + if (range.slice_divisor != 1) + { + // floor(n/d) from a cold-computed reciprocal; the low estimate is at + // most one short. No runtime integer division in the TMA issue loop. + const uint64_t estimate = __umul64hi(slice, range.slice_reciprocal); + slice = estimate + (slice - estimate * range.slice_divisor >= range.slice_divisor); + } + const uint64_t begin = slice * kSlice; + const uint64_t remaining = range.bytes - begin; + const uint64_t end = begin + (remaining < kSlice ? remaining : kSlice); + const uint64_t segment_end = range.segment_begin + range.segment_count; + lo = range.segment_begin; + hi = segment_end; + while (lo < hi) + { + const uint64_t mid = lo + (hi - lo) / 2; + if (segments[mid].virtual_begin + segments[mid].bytes <= begin) + lo = mid + 1; + else + hi = mid; + } + return Slice{begin, end, lo, segment_end}; +} + +__device__ __forceinline__ Fragment locate_fragment(Slice const& slice, MegamoeTmaCopySegment const& segment) +{ + const uint64_t begin = slice.begin > segment.virtual_begin ? slice.begin : segment.virtual_begin; + const uint64_t segment_end = segment.virtual_begin + segment.bytes; + const uint64_t end = slice.end < segment_end ? slice.end : segment_end; + unsigned const bytes = static_cast(end - begin); + const uint64_t offset = begin - segment.virtual_begin; + const uint64_t src = segment.src + offset, dst = segment.dst + offset; + // UC -> SMEM alignment depends only on the source. A differently aligned + // MC destination can consume the staged data with exact 32-bit stores. + unsigned prefix = static_cast((16U - (src & 15U)) & 15U); + if (prefix > bytes) + prefix = bytes; + unsigned const bulk = (bytes - prefix) & ~15U; + // The virtual concat is NOT a physical address space. Align each bulk + // fragment's scratch area independently; a preceding 4/8/12-byte plane + // must not force the next large weight onto the scalar path. + // end_of_bulk <= align_up(fragment_virtual_end,16) <= 8192, and the + // next fragment's aligned start is no smaller. Scratch regions are disjoint. + unsigned const smem_offset = static_cast((begin - slice.begin + prefix + 15U) & ~uint64_t(15U)); + return Fragment{src, dst, bytes, prefix, bulk, smem_offset}; +} + +// Only a virtual slice entirely contained in one physical segment can use +// one padded, aligned TMA load. Other slices retain exact scalar edge copies. +// The NOMINAL aligned 16-byte accesses, not just their masks, must lie within +// the original source and destination intervals. No small-plane overread. +__device__ __forceinline__ unsigned coalesced_copy_flags( + Slice const& slice, MegamoeTmaCopySegment const& segment, Fragment const& f) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + if (slice.begin < segment.virtual_begin || slice.end > segment.virtual_begin + segment.bytes + || ((f.src ^ f.dst) & 15U) || segment.bytes < 16) + return 0; + unsigned const tail = f.bytes - f.scalar_prefix - f.bulk_bytes; + unsigned const flags = (f.scalar_prefix ? 1U : 0U) | (tail ? 2U : 0U); + const uint64_t first_src = f.src & ~uint64_t(15); + const uint64_t first_dst = f.dst & ~uint64_t(15); + const uint64_t last_src = (f.src + f.bytes - 1) & ~uint64_t(15); + const uint64_t last_dst = (f.dst + f.bytes - 1) & ~uint64_t(15); + // Subtraction form also avoids overflow of a one-past/rounded-up pointer. + if (first_src < segment.src || first_dst < segment.dst || last_src < segment.src || last_dst < segment.dst + || last_src - segment.src > segment.bytes - 16 || last_dst - segment.dst > segment.bytes - 16) + return 0; + return flags | kCoalesced; +#else + return 0; +#endif +} + +__device__ __forceinline__ uint16_t edge_byte_mask(unsigned begin, unsigned bytes) +{ + // begin+bytes <= 16; the shift is evaluated in 32 bits. + return static_cast(((1U << bytes) - 1U) << begin); +} + +__device__ __forceinline__ void prefetch_bank(MegamoeTmaCopySegment const* segments, MegamoeTmaCopyRange const* ranges, + uint64_t range_count, uint64_t ordinal_base, uint64_t ordinal_stride, uint64_t sequence, unsigned char* warp_buffer, + SlotControl* warp_control, int slot_start, int count) +{ + for (int i = 0; i < count; ++i) + { + const Slice slice = locate_slice(segments, ranges, range_count, ordinal_base + ordinal_stride * (sequence + i)); + SlotControl* slot = warp_control + slot_start + i; + unsigned char* slot_buffer = warp_buffer + static_cast(slot_start + i) * kSlotDataBytes; + const MegamoeTmaCopySegment first_segment = segments[slice.segment_begin]; + const Fragment first = locate_fragment(slice, first_segment); + unsigned const fast = coalesced_copy_flags(slice, first_segment, first); + if (fast) + { + // ONE aligned TMA load supplies head, body and tail. The only extra + // storage is the final 16B of this physical slot; logical RR remains 8KiB. + slot->reserved = fast; + slot->fast_dst = first.dst; + slot->fast_bytes = first.bytes; + unsigned const load_bytes = (static_cast(first.src & 15U) + first.bytes + 15U) & ~15U; + expect_slice(slot, load_bytes); + load_fragment(slot_buffer, slot, first.src & ~uint64_t(15), load_bytes); + continue; + } + unsigned bulk_bytes = first.bulk_bytes; + for (uint64_t index = slice.segment_begin + 1; + index < slice.segment_end && segments[index].virtual_begin < slice.end; ++index) + bulk_bytes += locate_fragment(slice, segments[index]).bulk_bytes; + slot->reserved = bulk_bytes ? kFallbackBulk : 0U; + if (!bulk_bytes) + continue; + expect_slice(slot, bulk_bytes); + if (first.bulk_bytes) + load_fragment(slot_buffer + first.smem_offset, slot, first.src + first.scalar_prefix, first.bulk_bytes); + for (uint64_t index = slice.segment_begin + 1; + index < slice.segment_end && segments[index].virtual_begin < slice.end; ++index) + { + const Fragment f = locate_fragment(slice, segments[index]); + if (f.bulk_bytes) + load_fragment(slot_buffer + f.smem_offset, slot, f.src + f.scalar_prefix, f.bulk_bytes); + } + } +} + +__device__ __forceinline__ void scalar_words(Fragment const& f, unsigned begin, unsigned end) +{ + for (unsigned off = begin; off < end; off += 4) + { + unsigned value; + // UC means unicast. .cg avoids a stale read-only (.nc) cache. + asm volatile("ld.global.cg.u32 %0, [%1];" : "=r"(value) : "l"(f.src + off) : "memory"); + asm volatile("multimem.st.relaxed.sys.global.u32 [%0], %1;" ::"l"(f.dst + off), "r"(value) : "memory"); + } +} + +__device__ __forceinline__ void multicast_shared_words(Fragment const& f, unsigned char const* buffer) +{ + // The slot mbarrier has completed: TMA data is visible to generic shared + // loads. Only the source-aligned body is staged; no rounded overread. + for (unsigned off = 0; off < f.bulk_bytes; off += 4) + { + unsigned value; + asm volatile("ld.shared.u32 %0, [%1];" : "=r"(value) : "r"(shared_address(buffer + off)) : "memory"); + asm volatile("multimem.st.relaxed.sys.global.u32 [%0], %1;" ::"l"(f.dst + f.scalar_prefix + off), "r"(value) + : "memory"); + } +} + +__device__ __forceinline__ void consume_bank(MegamoeTmaCopySegment const* segments, MegamoeTmaCopyRange const* ranges, + uint64_t range_count, uint64_t ordinal_base, uint64_t ordinal_stride, uint64_t sequence, unsigned char* warp_buffer, + SlotControl* warp_control, int slot_start, int count, int bank, bool* pending, int& outstanding) +{ + bool issued = false; + for (int i = 0; i < count; ++i) + { + SlotControl* slot = warp_control + slot_start + i; + unsigned const flags = slot->reserved; + unsigned char const* slot_buffer = warp_buffer + static_cast(slot_start + i) * kSlotDataBytes; + if (flags & kCoalesced) + { + // Prefetch cached these values before issuing the TMA load. + // No range search, segment search, or descriptor reload on consume. + const uint64_t dst = slot->fast_dst; + unsigned const bytes = slot->fast_bytes; + unsigned const shift = static_cast(dst & 15U); + unsigned prefix = (16U - shift) & 15U; + if (prefix > bytes) + prefix = bytes; + unsigned const bulk = (bytes - prefix) & ~15U; + unsigned const tail_offset = prefix + bulk; + const uint16_t head_mask = edge_byte_mask(shift, prefix); + const uint16_t tail_mask = edge_byte_mask(0, bytes - tail_offset); + await_slice(slot); + if (bulk) + multicast_slice(dst + prefix, slot_buffer + shift + prefix, bulk); + if (flags & 1U) + multicast_edge(dst & ~uint64_t(15), slot_buffer, head_mask); + if (flags & 2U) + multicast_edge(dst + tail_offset, slot_buffer + shift + tail_offset, tail_mask); + issued = true; + continue; + } + const Slice slice = locate_slice(segments, ranges, range_count, ordinal_base + ordinal_stride * (sequence + i)); + // Prefetch already determined whether this slice has any async data: + // no duplicate fragment scan or second range lookup for scalar edges. + if (flags & kFallbackBulk) + await_slice(slot); + for (uint64_t index = slice.segment_begin; + index < slice.segment_end && segments[index].virtual_begin < slice.end; ++index) + { + const Fragment f = locate_fragment(slice, segments[index]); + if (f.bulk_bytes) + { + if (((f.dst + f.scalar_prefix) & 15U) == 0) + { + multicast_slice(f.dst + f.scalar_prefix, slot_buffer + f.smem_offset, f.bulk_bytes); + issued = true; + } + else + { + // Different UC/MC alignment cannot use a bulk MC store. + // Reuse the prefetched bytes instead of serial remote loads. + multicast_shared_words(f, slot_buffer + f.smem_offset); + } + } + // These exact scalar intervals are disjoint from ALL TMA writes. + // No intermediate group drain; final full-write wait/fence is unchanged. + scalar_words(f, 0, f.scalar_prefix); + scalar_words(f, f.scalar_prefix + f.bulk_bytes, f.bytes); + } + } + if (issued) + { + asm volatile("cp.async.bulk.commit_group;" ::: "memory"); + pending[bank] = true; + ++outstanding; + } +} + +// The two peers are the only writers of a bank ticket. Unsigned wrap is +// safe: a peer cannot publish its next ticket until the other responds. +__device__ __forceinline__ void wait_bank_ticket(unsigned const* address, unsigned expected) +{ + unsigned observed; + do + { + asm volatile("ld.acquire.cta.shared::cta.u32 %0, [%1];" + : "=r"(observed) + : "r"(shared_address(address)) + : "memory"); + if (observed != expected) + __nanosleep(32); + } while (observed != expected); +} + +__device__ __forceinline__ void publish_bank_ticket(unsigned* address, unsigned value) +{ + asm volatile("st.release.cta.shared::cta.u32 [%0], %1;" ::"r"(shared_address(address)), "r"(value) : "memory"); +} + +__device__ __forceinline__ void tma_copy_body(MegamoeTmaCopySegment const* segments, MegamoeTmaCopyRange const* ranges, + uint64_t range_count, uint64_t total_slices, int warps, int total_slots, unsigned* completed_ctas, uint64_t flag_mc, + uint64_t generation) +{ + if (!total_slices) + { + // Empty source ranks still publish READY. One tiny CTA, same kernel; + // no descriptor upload or additional host-side notification launch. + if (threadIdx.x == 0 && flag_mc) + multicast_generation(flag_mc, generation); + return; + } + extern __shared__ __align__(128) unsigned char smem[]; + int const physical_warp = threadIdx.x / 32, lane = threadIdx.x % 32; + int const warp = physical_warp / 2; // logical RR worker, shared by its two warps + bool const producer = (physical_warp & 1) == 0; + int const slots_per_warp = megamoe_tma_warp_slot_count(total_slots, warps, warp); + int const first_slot = megamoe_tma_warp_slot_begin(total_slots, warps, warp); + SlotControl* control = reinterpret_cast(smem + static_cast(total_slots) * kSlotDataBytes); + SlotControl* warp_control = control + first_slot; + unsigned char* warp_buffer = smem + static_cast(first_slot) * kSlotDataBytes; + if (producer && lane == 0) + { + for (int slot = 0; slot < slots_per_warp; ++slot) + initialize_barrier(warp_control + slot); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + } + __syncthreads(); + // Elect once for the entire pipeline so ptxas can prove one TMA issuer per warp. + if (cuda::ptx::elect_sync(0xffffffffu)) + { + // CTA/SM is the fastest-varying dimension, then warp, then sequence. + const uint64_t ordinal_base = megamoe_tma_worker_first_slice(blockIdx.x, gridDim.x, warp); + const uint64_t ordinal_stride = megamoe_tma_worker_slice_stride(gridDim.x, warps); + const uint64_t tasks = ordinal_base < total_slices ? 1 + (total_slices - 1 - ordinal_base) / ordinal_stride : 0; + int const capacity[2] = {(slots_per_warp + 1) / 2, slots_per_warp / 2}; + int const start[2] = {0, capacity[0]}; + unsigned ticket[2] = {producer ? 0U : 1U, producer ? 0U : 1U}; + uint64_t next_sequence = 0; + if (producer) + { + // Each bank is independent: bank 1 can already be loading while + // the consumer drains bank 0's shared-memory reads and releases it. + for (int bank = 0; next_sequence < tasks; bank ^= 1) + { + const uint64_t left = tasks - next_sequence; + int const count + = static_cast(left < static_cast(capacity[bank]) ? left : capacity[bank]); + unsigned* handshake = &warp_control[start[bank]].fast_padding; + wait_bank_ticket(handshake, ticket[bank]); + prefetch_bank(segments, ranges, range_count, ordinal_base, ordinal_stride, next_sequence, warp_buffer, + warp_control, start[bank], count); + // Publish generic slot metadata after all TMA loads have been + // issued. The consumer still waits on each load's mbarrier. + publish_bank_ticket(handshake, ticket[bank] + 1U); + ticket[bank] += 2U; + next_sequence += count; + } + } + else + { + bool pending[2] = {false, false}; + int outstanding = 0; + int held_bank = -1; + for (int bank = 0; next_sequence < tasks; bank ^= 1) + { + const uint64_t left = tasks - next_sequence; + int const count + = static_cast(left < static_cast(capacity[bank]) ? left : capacity[bank]); + unsigned* handshake = &warp_control[start[bank]].fast_padding; + wait_bank_ticket(handshake, ticket[bank]); + consume_bank(segments, ranges, range_count, ordinal_base, ordinal_stride, next_sequence, warp_buffer, + warp_control, start[bank], count, bank, pending, outstanding); + if (held_bank >= 0) + { + // Consume both banks before releasing the older one. The + // producer refills it while this bank's MC reads remain + // outstanding; at most two read-pending groups exist. + if (pending[held_bank]) + { + if (outstanding == 2) + wait_older_reads(); + else + wait_all_reads(); // this bank was scalar-only + pending[held_bank] = false; + --outstanding; + } + publish_bank_ticket(&warp_control[start[held_bank]].fast_padding, ticket[held_bank] + 1U); + ticket[held_bank] += 2U; + } + held_bank = bank; + next_sequence += count; + } + // The final bank has no following consume to release it. Flush + // remaining reads and publish its final free ticket explicitly. + if (held_bank >= 0) + { + if (pending[held_bank]) + wait_all_reads(); + publish_bank_ticket(&warp_control[start[held_bank]].fast_padding, ticket[held_bank] + 1U); + } + // READY still covers destination completion, not merely SMEM reads. + wait_all_writes(); + // Every issuer joins through the CTA barrier and the acq_rel + // counter chain before the final CTA publishes READY. + } + } + __syncthreads(); + if (threadIdx.x == 0 && flag_mc) + finish_cta_and_notify(completed_ctas, flag_mc, generation); +} + +// Per-CTA scratch is allocated once. No plan tensor is copied between kernels. +struct GpuDirectStorage +{ + MegamoeTmaGpuPlanConfig config; + int const *ids, *levels, *owners, *workspace, *status, *grid_sync; + int capacity; + uint64_t max_segments; + int* scratch; + MegamoeTmaCopySegment* segments; + MegamoeTmaCopyRange* ranges; + MegamoeTmaGpuPlanResult* results; +}; + +__global__ void tma_copy_gpu_direct_kernel(GpuDirectStorage direct, int warps, int total_slots, + unsigned* completed_ctas, uint64_t flag_mc, uint64_t generation) +{ + auto const& c = direct.config; + const uint64_t index = blockIdx.x; + int* plan = direct.scratch + index * (c.plan_words + 4ULL * c.helper_count); + auto* segments = direct.segments + index * direct.max_segments; + auto* ranges = direct.ranges + index * direct.max_segments; + auto* result = direct.results + index; + __shared__ int schedulerPlanReady; + if (threadIdx.x == 0) + { + bool statusOk = true; + for (int field = 0; field < 6; ++field) + statusOk &= direct.status[field] == 0; + auto const planEpoch = static_cast(direct.workspace[0]); + auto const schedulerEpoch = static_cast(direct.grid_sync[1]); + schedulerPlanReady = statusOk && planEpoch == schedulerEpoch; + if (!schedulerPlanReady && blockIdx.x == 0) + printf("GPU-direct in-switch copy: stale or failed scheduler plan; no READY published\n"); + } + __syncthreads(); + if (!schedulerPlanReady) + { + asm volatile("trap;" ::: "memory"); + return; + } + // HALO-Q's existing workspace: epoch, count, six capacity-strided fields. + int const* fields = direct.workspace + 2; + int const cap = direct.capacity; + tensorrt_llm::kernels::sami_copy::PlacementView placement{direct.workspace[1], cap, fields, fields + cap, + fields + 2 * cap, fields + 3 * cap, fields + 4 * cap, fields + 5 * cap}; + // Private routing scratch, not a CPU channel: no SYS publication, extra + // CUDA operation or scheduler-side copy of the placement outputs. + tensorrt_llm::kernels::sami_copy::publish_plan_channel_warp(plan, direct.ids, direct.levels, direct.owners, + placement, c.helper_count, c.world, c.rank, c.global_experts, c.plan_abi_version, c.route_features); + __syncthreads(); + if (threadIdx.x == 0) + { + megamoe_tma_build_gpu_plan( + &c, plan, plan + c.plan_words, segments, direct.max_segments, ranges, direct.max_segments, result); + // Malformed plans must never publish READY. The CUDA error surfaces at + // the caller's normal completion check, without a host wait in submit. + if (result->error) + { + printf("GPU-direct in-switch copy: invalid plan, error=%d; no READY published\n", result->error); + asm volatile("trap;" ::: "memory"); + } + } + __syncthreads(); + if (!result->total_slices) + { + // Work size is device-derived, so the launch always has config.sms CTAs. + // Empty ranks must still join the last-CTA chain before one READY store. + if (threadIdx.x == 0) + finish_cta_and_notify(completed_ctas, flag_mc, generation); + return; + } + tma_copy_body(segments, ranges, result->range_count, result->total_slices, warps, total_slots, completed_ctas, + flag_mc, generation); +} + +} // namespace + +struct MegamoeTmaCopyState +{ + MegamoeTmaCopyConfig config{}; + unsigned* completed_ctas = nullptr; + cudaStream_t last_stream = nullptr; + bool work_queued = false; + cudaError_t poison = cudaSuccess; + GpuDirectStorage direct{}; + uint64_t* direct_tables = nullptr; + bool direct_configured = false; + bool direct_bound = false; + std::mutex submit_mutex; +}; + +extern "C" int megamoe_tma_copy_destroy(MegamoeTmaCopyState** state_ptr) +{ + if (!state_ptr || !*state_ptr) + return static_cast(cudaSuccess); + MegamoeTmaCopyState* state = *state_ptr; + cudaError_t first = cudaSuccess; + auto record = [&first](cudaError_t error) + { + if (first == cudaSuccess) + first = error; + }; + int previous_device = -1; + cudaError_t error = cudaGetDevice(&previous_device); + record(error); + if (error == cudaSuccess && previous_device != state->config.device) + { + error = cudaSetDevice(state->config.device); + record(error); + } + if (error != cudaSuccess) + return static_cast(first); // Preserve state for a later cleanup attempt. + // GPU-direct submissions use one stream; synchronize it before releasing + // the bound plan storage and completion counter. + if (state->work_queued) + record(cudaStreamSynchronize(state->last_stream)); + if (state->direct_tables) + record(cudaFree(state->direct_tables)); + if (state->direct.scratch) + record(cudaFree(state->direct.scratch)); + if (state->direct.segments) + record(cudaFree(state->direct.segments)); + if (state->direct.ranges) + record(cudaFree(state->direct.ranges)); + if (state->direct.results) + record(cudaFree(state->direct.results)); + if (state->completed_ctas) + record(cudaFree(state->completed_ctas)); + if (previous_device != state->config.device) + record(cudaSetDevice(previous_device)); + delete state; + *state_ptr = nullptr; + return static_cast(first); +} + +extern "C" int megamoe_tma_copy_create(uint64_t max_segments, int sms, int warps, MegamoeTmaCopyState** out) +{ + if (!out) + return static_cast(cudaErrorInvalidValue); + *out = nullptr; + if (!max_segments || max_segments > SIZE_MAX / (sizeof(MegamoeTmaCopySegment) + sizeof(MegamoeTmaCopyRange)) + || sms < 0 || sms > MEGAMOE_TMA_COPY_MAX_SMS || warps < 0 || warps > MEGAMOE_TMA_COPY_MAX_WARPS) + return static_cast(cudaErrorInvalidValue); + auto* state = new (std::nothrow) MegamoeTmaCopyState; + if (!state) + return static_cast(cudaErrorMemoryAllocation); + auto fail = [&state](cudaError_t error) + { + megamoe_tma_copy_destroy(&state); + return static_cast(error); + }; + cudaError_t error = cudaGetDevice(&state->config.device); + if (error != cudaSuccess) + { + delete state; + return static_cast(error); + } + cudaDeviceProp prop{}; + error = cudaGetDeviceProperties(&prop, state->config.device); + if (error != cudaSuccess) + return fail(error); + int runtime_version = 0, multicast_supported = 0; + error = cudaRuntimeGetVersion(&runtime_version); + if (error != cudaSuccess) + return fail(error); + if (runtime_version < 13010 || prop.major < 9) + return fail(cudaErrorNotSupported); + // Multicast support is exposed by the Driver API, not cudaDeviceAttr. + CUdevice driver_device; + CUresult driver_error = cuDeviceGet(&driver_device, state->config.device); + if (driver_error != CUDA_SUCCESS) + return fail(cudaErrorUnknown); + driver_error = cuDeviceGetAttribute(&multicast_supported, CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED, driver_device); + if (driver_error != CUDA_SUCCESS) + return fail(driver_error == CUDA_ERROR_NOT_SUPPORTED ? cudaErrorNotSupported : cudaErrorUnknown); + if (!multicast_supported) + return fail(cudaErrorNotSupported); + auto& config = state->config; + config.abi_version = MEGAMOE_TMA_COPY_ABI_VERSION; + config.device_sm_count = prop.multiProcessorCount; + config.sms = megamoe_tma_copy_sm_count(sms, prop.multiProcessorCount); + config.warps = warps ? warps : 7; + config.threads_per_cta = config.warps * MEGAMOE_TMA_COPY_THREADS_PER_WORKER; + config.slice_bytes = kSlice; + config.max_segments = max_segments; + config.compute_major = prop.major; + config.compute_minor = prop.minor; + if (!config.sms) + return fail(cudaErrorInvalidValue); + error = cudaDeviceGetAttribute( + &config.device_optin_shared_bytes, cudaDevAttrMaxSharedMemoryPerBlockOptin, config.device); + if (error != cudaSuccess) + return fail(error); + error = cudaDeviceGetAttribute( + &config.device_shared_bytes_per_sm, cudaDevAttrMaxSharedMemoryPerMultiprocessor, config.device); + if (error != cudaSuccess) + return fail(error); + cudaFuncAttributes attributes{}; + error = cudaFuncGetAttributes(&attributes, tma_copy_gpu_direct_kernel); + if (error != cudaSuccess) + return fail(error); + const size_t available = config.device_optin_shared_bytes > attributes.sharedSizeBytes + ? config.device_optin_shared_bytes - attributes.sharedSizeBytes + : 0; + config.total_slots = megamoe_tma_copy_total_slots(static_cast(available)); + config.max_warps = megamoe_tma_copy_max_warps( + config.total_slots, std::min(prop.maxThreadsPerBlock, attributes.maxThreadsPerBlock)); + if (config.warps > config.max_warps) + return fail(cudaErrorInvalidConfiguration); + const size_t raw = static_cast(config.total_slots) * (kSlotDataBytes + kSlotControlBytes); + const size_t dynamic_bytes = (raw + 127) & ~size_t(127); + if (dynamic_bytes + attributes.sharedSizeBytes <= static_cast(config.device_shared_bytes_per_sm) / 2) + return fail(cudaErrorNotSupported); + config.slots_per_warp = config.total_slots / config.warps; + config.extra_slot_warps = config.total_slots % config.warps; + config.max_slots_per_warp = config.slots_per_warp + (config.extra_slot_warps != 0); + config.bank0_slots_per_warp = (config.slots_per_warp + 1) / 2; + config.bank1_slots_per_warp = config.slots_per_warp / 2; + config.dynamic_shared_bytes = static_cast(dynamic_bytes); + error = cudaFuncSetAttribute( + tma_copy_gpu_direct_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, config.dynamic_shared_bytes); + if (error != cudaSuccess) + return fail(error); + error = cudaFuncSetAttribute( + tma_copy_gpu_direct_kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared); + if (error != cudaSuccess) + return fail(error); + error = cudaOccupancyMaxActiveBlocksPerMultiprocessor(&config.max_active_ctas_per_sm, tma_copy_gpu_direct_kernel, + config.threads_per_cta, config.dynamic_shared_bytes); + if (error != cudaSuccess) + return fail(error); + if (config.max_active_ctas_per_sm != 1) + return fail(cudaErrorNotSupported); + error = cudaMalloc(reinterpret_cast(&state->completed_ctas), sizeof(unsigned)); + if (error != cudaSuccess) + return fail(error); + error = cudaMemset(state->completed_ctas, 0, sizeof(unsigned)); + if (error != cudaSuccess) + return fail(error); + // Cold initialization must be visible even when the first launch uses a + // non-default stream. Subsequent launches rearm entirely on the GPU. + error = cudaStreamSynchronize(nullptr); + if (error != cudaSuccess) + return fail(error); + *out = state; + return static_cast(cudaSuccess); +} + +extern "C" int megamoe_tma_copy_config_info(MegamoeTmaCopyState const* state, MegamoeTmaCopyConfig* out) +{ + if (!state || !out) + return static_cast(cudaErrorInvalidValue); + *out = state->config; + return static_cast(cudaSuccess); +} + +extern "C" char const* megamoe_tma_copy_error_string(int error) +{ + if (error == static_cast(cudaErrorInvalidConfiguration)) + return "invalid TMA launch geometry: each communication warp requires at least " + "two 8 KiB slices; shared-memory capacity and kernel/device thread limits apply"; + return cudaGetErrorString(static_cast(error)); +} + +extern "C" int megamoe_tma_copy_configure_gpu_plan(MegamoeTmaCopyState* state, MegamoeTmaGpuPlanConfig const* input) +{ + if (!state || !input) + return cudaErrorInvalidValue; + std::lock_guard lock(state->submit_mutex); + if (state->work_queued || state->direct_configured || state->direct_tables) + return cudaErrorInvalidValue; + auto& d = state->direct; + auto const& c = *input; + if (c.world < 2 || c.world > 32 || (c.world & 1) || c.rank < 0 || c.rank >= c.world || c.global_experts % c.world + || c.home_count != c.global_experts / c.world || c.plan_abi_version != 6 + || (c.route_features != 0 && c.route_features != 1) || c.tma_route < 0 || c.tma_route > 2 + || c.tma_source_load_percent < 1 || c.tma_source_load_percent > 199 || c.helper_count <= 0 + || c.helper_count > (INT32_MAX - 25) / 7 || c.global_experts <= 0 || c.global_experts > 384 || c.planes <= 0 + || c.planes > 32 || c.level_count <= 0 || c.level_count > 4 || c.target_count < 0 || c.target_count > 32 + || !c.src_table || !c.dst_table || !c.plane_bytes || (c.route_features && !c.foreign_dst_table)) + return cudaErrorInvalidValue; + if (c.owner_stride != ((c.helper_count + 3) & ~3) || c.plan_words != 4 + 7 * c.owner_stride) + return cudaErrorInvalidValue; + int expected_targets = 0; + for (int level = 0; level < c.level_count; ++level) + { + int group = c.group_sizes[level]; + if (group < (level ? 4 : 2) || group > c.world || c.world % group || (!level && group != c.world) + || (level && (group >= c.group_sizes[level - 1] || c.group_sizes[level - 1] % group))) + return cudaErrorInvalidValue; + if (level) + expected_targets += c.world / group; + } + if (c.target_count != expected_targets || (c.route_features && !expected_targets) + || (!c.route_features && c.foreign_dst_table)) + return cudaErrorInvalidValue; + uint64_t total_bytes = 0; + for (int plane = 0; plane < c.planes; ++plane) + { + const uint64_t bytes = c.plane_bytes[plane]; + if (!bytes || (bytes & 3) || bytes > UINT64_MAX / c.world - total_bytes) + return cudaErrorInvalidValue; + total_bytes += bytes; + } + auto valid_address = [&](uint64_t address, int plane) + { return address && !(address & 3) && c.plane_bytes[plane] <= UINT64_MAX - address; }; + for (uint64_t i = 0; i < uint64_t(c.global_experts) * c.planes; ++i) + if (!valid_address(c.src_table[i], i % c.planes)) + return cudaErrorInvalidValue; + for (uint64_t i = 0; i < uint64_t(c.level_count) * c.helper_count * c.planes; ++i) + if (!valid_address(c.dst_table[i], i % c.planes)) + return cudaErrorInvalidValue; + if (c.route_features) + { + uint64_t index = 0; + for (int level = 1; level < c.level_count; ++level) + { + int group = c.group_sizes[level]; + for (int begin = 0; begin < c.world; begin += group) + for (int slot = 0; slot < c.helper_count; ++slot) + for (int plane = 0; plane < c.planes; ++plane, ++index) + { + const uint64_t address = c.foreign_dst_table[index]; + bool const member = c.rank >= begin && c.rank < begin + group; + if (member ? address != 0 : !valid_address(address, plane)) + return cudaErrorInvalidValue; + } + } + } + int device = -1; + cudaError_t error = cudaGetDevice(&device); + if (error != cudaSuccess) + return error; + if (device != state->config.device) + return cudaErrorInvalidDevice; + auto fail = [state](cudaError_t e) + { + state->poison = e; + return static_cast(e); + }; + d.config = c; + const uint64_t sources = uint64_t(c.global_experts) * c.planes; + const uint64_t destinations = uint64_t(c.level_count) * c.helper_count * c.planes; + const uint64_t foreign = c.route_features ? uint64_t(c.target_count) * c.helper_count * c.planes : 0; + const uint64_t words = sources + destinations + foreign + c.planes; + if (words > SIZE_MAX / sizeof(uint64_t)) + return cudaErrorInvalidValue; + error = cudaMalloc(reinterpret_cast(&state->direct_tables), words * sizeof(uint64_t)); + if (error != cudaSuccess) + return fail(error); + uint64_t* cursor = state->direct_tables; + auto upload = [&](uint64_t const* host, uint64_t count, uint64_t const** target) + { + *target = count ? cursor : nullptr; + auto e = count ? cudaMemcpy(cursor, host, count * sizeof(uint64_t), cudaMemcpyHostToDevice) : cudaSuccess; + cursor += count; + return e; + }; + if ((error = upload(c.src_table, sources, &d.config.src_table)) != cudaSuccess + || (error = upload(c.dst_table, destinations, &d.config.dst_table)) != cudaSuccess + || (error = upload(c.foreign_dst_table, foreign, &d.config.foreign_dst_table)) != cudaSuccess + || (error = upload(c.plane_bytes, c.planes, &d.config.plane_bytes)) != cudaSuccess) + return fail(error); + auto const& config = state->config; + const uint64_t n = config.sms; + d.max_segments = config.max_segments; + if (d.max_segments > SIZE_MAX / n / sizeof(MegamoeTmaCopyRange)) + return fail(cudaErrorInvalidValue); + error = cudaMalloc(reinterpret_cast(&d.scratch), n * (c.plan_words + 4ULL * c.helper_count) * sizeof(int)); + if (error != cudaSuccess) + return fail(error); + error = cudaMalloc(reinterpret_cast(&d.segments), n * d.max_segments * sizeof(MegamoeTmaCopySegment)); + if (error != cudaSuccess) + return fail(error); + error = cudaMalloc(reinterpret_cast(&d.ranges), n * d.max_segments * sizeof(MegamoeTmaCopyRange)); + if (error != cudaSuccess) + return fail(error); + error = cudaMalloc(reinterpret_cast(&d.results), n * sizeof(MegamoeTmaGpuPlanResult)); + if (error != cudaSuccess) + return fail(error); + error = cudaMemset(d.results, 0, n * sizeof(MegamoeTmaGpuPlanResult)); + if (error != cudaSuccess) + return fail(error); + error = cudaStreamSynchronize(nullptr); // Cold initialization only. + if (error != cudaSuccess) + return fail(error); + state->direct_configured = true; + return cudaSuccess; +} + +extern "C" int megamoe_tma_copy_bind_gpu_direct(MegamoeTmaCopyState* state, uint64_t ids, uint64_t levels, + uint64_t owners, uint64_t workspace, uint64_t status, uint64_t grid_sync, int capacity) +{ + if (!state || !ids || !levels || !owners || !workspace || !status || !grid_sync + || (ids | levels | owners | workspace | status | grid_sync) % alignof(int) || capacity < 1 || capacity > 384) + return cudaErrorInvalidValue; + std::lock_guard lock(state->submit_mutex); + if (!state->direct_configured || state->direct_bound || state->work_queued) + return cudaErrorInvalidValue; + int device = -1; + auto error = cudaGetDevice(&device); + if (error != cudaSuccess) + return error; + if (device != state->config.device) + return cudaErrorInvalidDevice; + for (auto ptr : {ids, levels, owners, workspace, status, grid_sync}) + { + cudaPointerAttributes attr{}; + error = cudaPointerGetAttributes(&attr, reinterpret_cast(ptr)); + if (error != cudaSuccess) + return error; + if (attr.type != cudaMemoryTypeDevice || attr.device != device) + return cudaErrorInvalidValue; + } + auto& d = state->direct; + d.ids = reinterpret_cast(ids); + d.levels = reinterpret_cast(levels); + d.owners = reinterpret_cast(owners); + d.workspace = reinterpret_cast(workspace); + d.status = reinterpret_cast(status); + d.grid_sync = reinterpret_cast(grid_sync); + d.capacity = capacity; + state->direct_bound = true; + return cudaSuccess; +} + +extern "C" int megamoe_tma_copy_submit_gpu_direct( + MegamoeTmaCopyState* state, uint64_t flag_mc, uint64_t generation, void* cuda_stream) +{ + if (!state || !flag_mc || (flag_mc & 7) || !generation) + return cudaErrorInvalidValue; + std::lock_guard lock(state->submit_mutex); + if (state->poison != cudaSuccess) + return state->poison; + if (!state->direct_bound) + return cudaErrorInvalidValue; + int device = -1; + cudaError_t error = cudaGetDevice(&device); + if (error != cudaSuccess) + return error; + if (device != state->config.device) + return cudaErrorInvalidDevice; + auto stream = reinterpret_cast(cuda_stream); + cudaStreamCaptureStatus capture; + error = cudaStreamIsCapturing(stream, &capture); + if (error != cudaSuccess) + return error; + if (capture != cudaStreamCaptureStatusNone) + return cudaErrorStreamCaptureUnsupported; + // Raw direct ABI also fixes one stream: protecting descriptor scratch on a + // different stream would not protect HALO-Q's already-enqueued output writes. + if (state->work_queued && stream != state->last_stream) + return cudaErrorInvalidValue; + state->last_stream = stream; + state->work_queued = true; + auto const& c = state->config; + // Keep this an ordinary same-stream launch: programmatic stream serialization + // could observe a stale matching modular epoch before HALO-Q publishes. + tma_copy_gpu_direct_kernel<<>>( + state->direct, c.warps, c.total_slots, state->completed_ctas, flag_mc, generation); + error = cudaGetLastError(); + if (error != cudaSuccess) + { + state->poison = error; + return error; + } + return cudaSuccess; +} + +extern "C" int megamoe_tma_copy_gpu_plan_result(MegamoeTmaCopyState* state, MegamoeTmaGpuPlanResult* result) +{ + if (!state || !result) + return cudaErrorInvalidValue; + std::lock_guard lock(state->submit_mutex); + if (!state->direct_configured) + return cudaErrorInvalidValue; + if (state->poison != cudaSuccess) + return state->poison; + int device = -1; + auto device_error = cudaGetDevice(&device); + if (device_error != cudaSuccess) + return device_error; + if (device != state->config.device) + return cudaErrorInvalidDevice; + if (state->work_queued) + { + auto error = cudaStreamSynchronize(state->last_stream); + if (error != cudaSuccess) + { + state->poison = error; + return error; + } + } + return cudaMemcpy(result, state->direct.results, sizeof(*result), cudaMemcpyDeviceToHost); +} diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h new file mode 100644 index 000000000000..d32029dd7551 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h @@ -0,0 +1,203 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +/* CUDA TMA multicast-copy backend. C ABI; no CUDA or Python headers required. */ +#ifndef MEGAMOE_TMA_COPY_H +#define MEGAMOE_TMA_COPY_H + +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlan.h" +#include + +#ifdef __cplusplus +extern "C" +{ +#endif + +#define MEGAMOE_TMA_COPY_ABI_VERSION 6 +#define MEGAMOE_TMA_COPY_SLICE_BYTES 8192 +#define MEGAMOE_TMA_COPY_SLOT_CONTROL_BYTES 32 +#define MEGAMOE_TMA_COPY_SLOT_DATA_BYTES 8208 +#define MEGAMOE_TMA_COPY_MAX_SMS 8 +#define MEGAMOE_TMA_COPY_MAX_WARPS 32 +/* Experimental candidate: warps counts logical producer/consumer worker pairs. */ +#define MEGAMOE_TMA_COPY_THREADS_PER_WORKER 64 + + /* Zero requests the default; zero return denotes invalid geometry. The entire + * payload/notification kernel stays within this per-GPU copy SM budget. */ + static inline int megamoe_tma_copy_sm_count(int requested_sms, int device_sms) + { + if (device_sms <= 0 || requested_sms < 0 || requested_sms > MEGAMOE_TMA_COPY_MAX_SMS + || requested_sms > device_sms) + return 0; + return requested_sms ? requested_sms + : (device_sms < MEGAMOE_TMA_COPY_MAX_SMS ? device_sms : MEGAMOE_TMA_COPY_MAX_SMS); + } + + /* Each logical slice remains 8192B; physical slots add 16B for one aligned + * load that straddles the logical slice boundary. + * Maximum whole slots after control records and 128-byte launch alignment. + * Capacity is independent of warp count; remainder slots are not discarded. */ + static inline int megamoe_tma_copy_total_slots(int available_bytes) + { + if (available_bytes <= 0) + return 0; + int const slot_bytes = MEGAMOE_TMA_COPY_SLOT_DATA_BYTES + MEGAMOE_TMA_COPY_SLOT_CONTROL_BYTES; + int slots = available_bytes / slot_bytes; + while (slots > 0 && (((uint64_t) slots * slot_bytes + 127) & ~(uint64_t) 127) > (uint64_t) available_bytes) + --slots; + return slots; + } + + static inline int megamoe_tma_copy_max_warps(int total_slots, int max_threads) + { + if (total_slots < 2 || max_threads < MEGAMOE_TMA_COPY_THREADS_PER_WORKER) + return 0; + int warps = total_slots / 2; + if (warps > MEGAMOE_TMA_COPY_MAX_WARPS) + warps = MEGAMOE_TMA_COPY_MAX_WARPS; + if (warps > max_threads / MEGAMOE_TMA_COPY_THREADS_PER_WORKER) + warps = max_threads / MEGAMOE_TMA_COPY_THREADS_PER_WORKER; + return warps; + } + +#if defined(__CUDACC__) +#define MEGAMOE_TMA_HD __host__ __device__ +#else +#define MEGAMOE_TMA_HD +#endif + /* Validated geometry has warps>=1 and total_slots>=2*warps. The first + * total_slots%warps warps own one extra slot; every slot has exactly one owner. */ + static inline MEGAMOE_TMA_HD int megamoe_tma_warp_slot_count(int total_slots, int warps, int warp) + { + return total_slots / warps + (warp < total_slots % warps); + } + + static inline MEGAMOE_TMA_HD int megamoe_tma_warp_slot_begin(int total_slots, int warps, int warp) + { + int extra = total_slots % warps; + return warp * (total_slots / warps) + (warp < extra ? warp : extra); + } + + /* Shared with CPU coverage checks: CTA/SM first, warp second, sequence last. */ + static inline MEGAMOE_TMA_HD uint64_t megamoe_tma_worker_first_slice(int cta, int sms, int warp) + { + return (uint64_t) cta + (uint64_t) sms * (uint64_t) warp; + } + + static inline MEGAMOE_TMA_HD uint64_t megamoe_tma_worker_slice_stride(int sms, int warps) + { + return (uint64_t) sms * (uint64_t) warps; + } + +#undef MEGAMOE_TMA_HD + + typedef struct MegamoeTmaCopyState MegamoeTmaCopyState; + + /* A physical interval within one logically concatenated byte stream. Entries + * in a range are contiguous in virtual space (no per-weight/plane rounding). + * Addresses/lengths are multiples of four; source is unicast, dst is multicast. + */ + typedef struct MegamoeTmaCopySegment + { + uint64_t src; + uint64_t dst; + uint64_t bytes; + uint64_t virtual_begin; + } MegamoeTmaCopySegment; + + /* A whole routing domain: source-owned slots, or one scatter destination team. + * Slice j spans [j*8192,min((j+1)*8192,bytes)) in its virtual byte stream and can + * cross any number of segments. It keeps ONE GPU/CTA/warp owner throughout. + * This rank owns floor((first_slice+j*slice_stride)/slice_divisor). Uniform + * routing uses divisor=1, reciprocal=0 and retains the original integer path. + * Otherwise reciprocal=UINT64_MAX/divisor enables an exact mulhi+correction. + * prefix_end compacts rank-local slices across ranges for SM/worker round-robin. + * No reset at weight boundaries. + */ + typedef struct MegamoeTmaCopyRange + { + uint64_t segment_begin; + uint64_t segment_count; + uint64_t bytes; + uint64_t first_slice; + uint64_t slice_stride; + uint64_t slice_count; + uint64_t prefix_end; + uint64_t slice_divisor; + uint64_t slice_reciprocal; + } MegamoeTmaCopyRange; + + typedef struct MegamoeTmaCopyConfig + { + int abi_version; + int device; + int device_sm_count; + int sms; + int warps; + int threads_per_cta; + int slots_per_warp; /* Minimum over all warps, floor(total_slots/warps). */ + int bank0_slots_per_warp; /* Minimum bank 0 capacity, ceil(slots_per_warp/2). */ + int bank1_slots_per_warp; /* Minimum bank 1 capacity, floor(slots_per_warp/2). */ + int slice_bytes; + int dynamic_shared_bytes; + int device_optin_shared_bytes; + int device_shared_bytes_per_sm; + int max_active_ctas_per_sm; + int compute_major; + int compute_minor; + int total_slots; + int max_slots_per_warp; + int extra_slot_warps; /* First this many warps receive max_slots_per_warp. */ + int max_warps; /* Shared-memory and kernel/device thread capacity. */ + uint64_t max_segments; + } MegamoeTmaCopyConfig; + + /* Every function returns a CUDA runtime error integer (0 == success). + * sms=0 selects min(device_sm_count,8); otherwise 1..min(device_sm_count,8). + * Explicit counts above eight are rejected; READY stays in this same kernel. + * warps=0 selects seven communication warps; explicit requests are 1..32. + * Runtime geometry requires at least two 8 KiB slots per warp and respects + * the active kernel and device limits. + * All whole shared-memory slots are divided near-equally among warps, reserving + * an mbarrier/phase record for every slot. Two banks use ALL allocated slots. + * Excessive runtime warp counts return cudaErrorInvalidConfiguration. + * Creation verifies that this shared-memory footprint permits one CTA/SM. + */ + int megamoe_tma_copy_create(uint64_t max_segments, int sms, int warps, MegamoeTmaCopyState** out); + + /* GPU-direct cold setup: input config contains HOST pointer tables. They are + * uploaded once; current HALO-Q outputs are bound by device address, never copied. + * Submit is stream-ordered and does not read or wait for plan data on the CPU. + * All bound tensors must remain alive and only be reused on the same stream + * after the copy kernel has read them. The Python endpoint enforces this lease. + */ + int megamoe_tma_copy_configure_gpu_plan(MegamoeTmaCopyState* state, MegamoeTmaGpuPlanConfig const* config); + int megamoe_tma_copy_bind_gpu_direct(MegamoeTmaCopyState* state, uint64_t ids, uint64_t levels, uint64_t owners, + uint64_t workspace, uint64_t status, uint64_t grid_sync, int capacity); + int megamoe_tma_copy_submit_gpu_direct( + MegamoeTmaCopyState* state, uint64_t flag_mc, uint64_t generation, void* cuda_stream); + /* Explicit diagnostic only: synchronizes last work and copies one small result. */ + int megamoe_tma_copy_gpu_plan_result(MegamoeTmaCopyState* state, MegamoeTmaGpuPlanResult* result); + /* Nulls *state once ownership is consumed. If device activation fails before + * cleanup begins, the pointer is preserved so the caller can retry. */ + int megamoe_tma_copy_destroy(MegamoeTmaCopyState** state); + int megamoe_tma_copy_config_info(MegamoeTmaCopyState const* state, MegamoeTmaCopyConfig* out); + char const* megamoe_tma_copy_error_string(int error); + +#ifdef __cplusplus +} +#endif +#endif diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlan.h b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlan.h new file mode 100644 index 000000000000..9ff8abd2eaf3 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlan.h @@ -0,0 +1,53 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +/* Device-resident ABI6 plan input for the fused TMA copy kernel. */ +#ifndef MEGAMOE_TMA_GPU_PLAN_H +#define MEGAMOE_TMA_GPU_PLAN_H +#include + +#define MEGAMOE_TMA_GPU_PLAN_MAX_WORLD 32 +#define MEGAMOE_TMA_GPU_PLAN_MAX_LEVELS 4 + +/* Immutable geometry and device addresses. Upload these tables once at bind. + * Source entries: [global_experts][planes]. Destinations: [level][helper][plane]. + * External destinations: [target][helper][plane]. Bytes: [plane]. No entry here + * depends on the current HALO-Q result. All pointer values are device pointers. + */ +typedef struct MegamoeTmaGpuPlanConfig +{ + int world, rank, helper_count, planes, global_experts, home_count; + int owner_stride, level_count, target_count; + int plan_abi_version, plan_words, route_features; + int tma_route; /* 0: published modes; 1: source; 2: scatter. */ + int tma_source_load_percent; /* Source quota = slices * percent / (100 * team). */ + int group_sizes[MEGAMOE_TMA_GPU_PLAN_MAX_LEVELS]; + uint64_t const *src_table, *dst_table, *foreign_dst_table, *plane_bytes; +} MegamoeTmaGpuPlanConfig; + +/* Per-CTA output; only successful construction makes counts consumable. + * Error 1: malformed plan/config; 2: forced source outside its destination team; + * 3: invalid address, byte count, capacity, or arithmetic overflow. + */ +typedef struct MegamoeTmaGpuPlanResult +{ + int error; + int direct_slots, scatter_slots; + int weighted_domains, uniform_domains, mixed_source_domains, foreign_source_domains; + uint64_t segment_count, range_count, total_slices; + uint64_t weighted_total_slices, weighted_source_quota_slices, weighted_local_slices; +} MegamoeTmaGpuPlanResult; +#endif diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlanBuilder.cuh b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlanBuilder.cuh new file mode 100644 index 000000000000..351bc18774b7 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlanBuilder.cuh @@ -0,0 +1,492 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +/* Construct the existing TMA descriptor ABI inside the payload kernel. + * Each CTA has private global scratch: 4*helper_count int32 slot words, + * max_segments segments, max_ranges ranges, and one result. CTA thread zero + * calls the builder and all threads synchronize before consuming its output. + * No CPU plan read, descriptor DMA, allocation, atomics or device-wide barrier. + * Scheduling already orders the producer before this kernel on its stream. + */ +#ifndef MEGAMOE_TMA_GPU_PLAN_BUILDER_CUH +#define MEGAMOE_TMA_GPU_PLAN_BUILDER_CUH +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h" +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlan.h" +#include +#if defined(__CUDACC__) +#define MEGAMOE_GP_HD __host__ __device__ +#else +#define MEGAMOE_GP_HD +#endif +namespace megamoe_gpu_plan_detail +{ +struct Builder : MegamoeTmaGpuPlanConfig +{ + int32_t const* plan; + int32_t* slot_scratch; + MegamoeTmaCopySegment* tma_segments; + MegamoeTmaCopyRange* tma_ranges; + size_t max_commands, tma_max_ranges; + size_t tma_segment_count, tma_range_count, last_commands; + int last_direct_slots, last_scatter_slots; + int last_tma_weighted_domains, last_tma_uniform_domains; + int last_tma_mixed_source_domains, last_tma_foreign_source_domains; + uint64_t last_tma_weighted_total_slices, last_tma_weighted_source_quota_slices; + uint64_t last_tma_weighted_local_slices, last_tma_local_slices; +}; + +static MEGAMOE_GP_HD inline uint64_t tma_weighted_source_quota(uint64_t slices, unsigned members, unsigned percent) +{ + return slices * percent / (100U * members); +} + +static MEGAMOE_GP_HD inline int target_geometry(Builder const* c, int target, int* level_out, int* begin_out) +{ + if (target < 0 || target >= c->target_count) + return -1; + for (int level = 1; level < c->level_count; ++level) + { + int count = c->world / c->group_sizes[level]; + if (target < count) + { + *level_out = level; + *begin_out = target * c->group_sizes[level]; + return 0; + } + target -= count; + } + return -1; +} + +static MEGAMOE_GP_HD inline int append_tma_segment( + Builder* c, size_t* segments, uint64_t* cursor, uint64_t src, uint64_t dst, size_t bytes) +{ + if (*segments >= c->max_commands || !src || !dst || !bytes || (src & 3) || (dst & 3) || (bytes & 3) + || bytes > UINT64_MAX - *cursor || bytes > UINT64_MAX - src || bytes > UINT64_MAX - dst) + return -1; + MegamoeTmaCopySegment* s = &c->tma_segments[(*segments)++]; + s->src = src; + s->dst = dst; + s->bytes = bytes; + s->virtual_begin = *cursor; + *cursor += bytes; + return 0; +} + +/* The validated arena binds one UC object per (owner, plane) and one MC + * object per (team, plane). The caller resets identity at plane/domain + * boundaries and gives normal and external teams separate namespaces. + * Numerical pointer adjacency alone is not enough to merge MC destinations. */ +static MEGAMOE_GP_HD inline int append_tma_coalesced_segment(Builder* c, size_t* segments, uint64_t* cursor, + uint64_t src, uint64_t dst, size_t bytes, uint64_t identity, uint64_t* previous_identity) +{ + if (!src || !dst || !bytes || (src & 3) || (dst & 3) || (bytes & 3) || bytes > UINT64_MAX - *cursor + || bytes > UINT64_MAX - src || bytes > UINT64_MAX - dst) + return -1; + if (*segments && *previous_identity == identity) + { + MegamoeTmaCopySegment* previous = &c->tma_segments[*segments - 1]; + if (previous->bytes <= UINT64_MAX - previous->src && previous->bytes <= UINT64_MAX - previous->dst + && previous->bytes <= UINT64_MAX - previous->virtual_begin && previous->src + previous->bytes == src + && previous->dst + previous->bytes == dst && previous->virtual_begin + previous->bytes == *cursor + && bytes <= UINT64_MAX - previous->bytes) + { + previous->bytes += bytes; + *cursor += bytes; + return 0; + } + } + if (append_tma_segment(c, segments, cursor, src, dst, bytes)) + return -1; + *previous_identity = identity; + return 0; +} + +static MEGAMOE_GP_HD inline int append_tma_range(Builder* c, size_t* ranges, uint64_t* prefix, size_t segment_begin, + size_t segment_count, uint64_t bytes, uint64_t first, uint64_t step, uint64_t slices, uint64_t divisor, + uint64_t reciprocal) +{ + /* These constant bounds also avoid runtime division for overflow checks. + * Valid 8KiB streams are smaller than this conservative slice bound. */ + if (!step || !divisor || divisor > 100U * MEGAMOE_TMA_GPU_PLAN_MAX_WORLD + || slices > UINT64_MAX / (100U * MEGAMOE_TMA_GPU_PLAN_MAX_WORLD)) + return -1; + uint64_t scaled_limit = slices * divisor; + if (first >= scaled_limit) + return 0; + uint64_t count = 1 + (scaled_limit - 1 - first) / step; + if (*ranges >= c->tma_max_ranges || count > UINT64_MAX - *prefix) + return -1; + MegamoeTmaCopyRange* r = &c->tma_ranges[(*ranges)++]; + r->segment_begin = segment_begin; + r->segment_count = segment_count; + r->bytes = bytes; + r->first_slice = first; + r->slice_stride = step; + r->slice_count = count; + *prefix += count; + r->prefix_end = *prefix; + r->slice_divisor = divisor; + r->slice_reciprocal = reciprocal; + return 0; +} + +static MEGAMOE_GP_HD inline size_t build_tma_descriptors(Builder* c, int outgoing_count, int* plan_error) +{ + *plan_error = 3; + int slots = c->helper_count, stride = c->owner_stride; + int *experts = c->slot_scratch, *levels = experts + slots; + int *owners = levels + slots, *modes = owners + slots; + size_t segments = 0, ranges = 0; + uint64_t prefix = 0; + int direct = 0, scatter = 0; + c->last_tma_weighted_domains = c->last_tma_uniform_domains = 0; + c->last_tma_mixed_source_domains = c->last_tma_foreign_source_domains = 0; + c->last_tma_weighted_total_slices = c->last_tma_weighted_source_quota_slices = 0; + c->last_tma_weighted_local_slices = c->last_tma_local_slices = 0; + for (int slot = 0; slot < slots; ++slot) + { + if (experts[slot] < 0 || modes[slot] == 2) + continue; + if (c->tma_route == 1) + { + int group = c->group_sizes[levels[slot]]; + int begin = c->rank / group * group; + if (owners[slot] < begin || owners[slot] >= begin + group) + { + *plan_error = 2; + return 0; + } + modes[slot] = 1; + } + else if (c->tma_route == 2) + { + modes[slot] = 0; + } + direct += modes[slot] == 1; + scatter += modes[slot] == 0; + } + /* Domain zero is source-owned work (including outgoing direct). A scatter + * domain must have one common member set, so it is separated per level; + * all ranks in that aligned team see the same plane/slot byte order. */ + for (int domain = 0; domain <= c->level_count; ++domain) + { + size_t begin_segment = segments; + uint64_t cursor = 0; + /* Keep one uninterrupted virtual stream per routing domain. Plane-major + * order groups the large aligned weights before the scalar planes; the + * physical arena and source ownership are unchanged. Within each plane, + * normal source-owned slots precede existing external outgoing records. */ + for (int plane = 0; plane < c->planes; ++plane) + { + uint64_t previous_identity = UINT64_MAX; + for (int slot = 0; slot < slots; ++slot) + { + if (experts[slot] < 0 || modes[slot] == 2) + continue; + if (domain == 0 ? (modes[slot] != 1 || owners[slot] != c->rank) + : (modes[slot] != 0 || levels[slot] != domain - 1)) + continue; + int level = levels[slot]; + size_t si = (size_t) experts[slot] * c->planes + plane; + size_t di = ((size_t) level * slots + slot) * c->planes + plane; + uint64_t identity = ((uint64_t) owners[slot] << 32) | (uint32_t) level; + if (append_tma_coalesced_segment(c, &segments, &cursor, c->src_table[si], c->dst_table[di], + c->plane_bytes[plane], identity, &previous_identity)) + return 0; + } + if (domain == 0) + { + for (int i = 0; i < outgoing_count; ++i) + { + int expert = c->plan[4 + 4 * stride + i]; + int target = c->plan[4 + 5 * stride + i]; + int helper = c->plan[4 + 6 * stride + i]; + size_t si = (size_t) expert * c->planes + plane; + size_t di = ((size_t) target * slots + helper) * c->planes + plane; + uint64_t identity = ((uint64_t) c->rank << 32) | (uint32_t) (c->level_count + target); + if (append_tma_coalesced_segment(c, &segments, &cursor, c->src_table[si], c->foreign_dst_table[di], + c->plane_bytes[plane], identity, &previous_identity)) + return 0; + } + } + } + if (!cursor) + continue; + uint64_t step = domain ? (uint64_t) c->group_sizes[domain - 1] : 1; + uint64_t first = domain ? (uint64_t) c->rank % step : 0; + uint64_t slices = cursor / MEGAMOE_TMA_COPY_SLICE_BYTES + (cursor % MEGAMOE_TMA_COPY_SLICE_BYTES != 0); + size_t before_ranges = ranges; + int common_source = -1; + int weighted = 0; + if (domain && c->tma_source_load_percent != 100) + { + for (int slot = 0; slot < slots; ++slot) + { + if (experts[slot] < 0 || modes[slot] != 0 || levels[slot] != domain - 1) + continue; + if (common_source == -1) + common_source = owners[slot]; + else if (common_source != owners[slot]) + { + common_source = -2; + break; + } + } + int team_begin = c->rank / (int) step * (int) step; + if (common_source == -2) + ++c->last_tma_mixed_source_domains; + else if (common_source < team_begin || common_source >= team_begin + (int) step) + ++c->last_tma_foreign_source_domains; + else + weighted = 1; + } + if (weighted) + { + unsigned members = (unsigned) step; + unsigned source_local = (unsigned) common_source % members; + unsigned percent = (unsigned) c->tma_source_load_percent; + uint64_t denominator = 100U * members; + uint64_t before_prefix = prefix; + uint64_t source_quota = tma_weighted_source_quota(slices, members, percent); + if (slices > UINT64_MAX - c->last_tma_weighted_total_slices + || source_quota > UINT64_MAX - c->last_tma_weighted_source_quota_slices) + return 0; + ++c->last_tma_weighted_domains; + c->last_tma_weighted_total_slices += slices; + c->last_tma_weighted_source_quota_slices += source_quota; + uint64_t divisor, reciprocal; + if (c->rank == common_source) + { + first = denominator - 1; + step = denominator; + divisor = percent; + reciprocal = divisor == 1 ? 0 : UINT64_MAX / divisor; + } + else + { + unsigned peer_order = ((unsigned) c->rank % members + members - source_local) % members - 1; + first = peer_order * denominator; + step = (members - 1) * denominator; + divisor = denominator - percent; + reciprocal = divisor == 1 ? 0 : UINT64_MAX / divisor; + } + /* floor((first+j*step)/divisor) is the inverse of the source + * quota events and peer round-robin. The whole existing virtual + * stream stays intact, with increasing slices and a continuous + * rank-local prefix across domains for SM/worker round-robin. */ + if (append_tma_range(c, &ranges, &prefix, begin_segment, segments - begin_segment, cursor, first, step, + slices, divisor, reciprocal)) + return 0; + c->last_tma_weighted_local_slices += prefix - before_prefix; + } + else + { + if (domain) + ++c->last_tma_uniform_domains; + if (append_tma_range( + c, &ranges, &prefix, begin_segment, segments - begin_segment, cursor, first, step, slices, 1, 0)) + return 0; + } + if (ranges == before_ranges) + segments = begin_segment; /* This rank owns none of this stream's slices. */ + } + c->last_tma_local_slices = prefix; + c->tma_segment_count = segments; + c->tma_range_count = ranges; + c->last_commands = segments; + c->last_direct_slots = direct; + c->last_scatter_slots = scatter; + *plan_error = 0; + return segments; +} + +static MEGAMOE_GP_HD inline size_t build_descriptors(Builder* c, int* plan_error) +{ + const volatile int32_t* plan = c->plan; + int const slots = c->helper_count; + int const stride = c->owner_stride; + int const quota = slots < c->home_count ? slots : c->home_count; + int* experts = c->slot_scratch; + int* levels = experts + slots; + int* owners = levels + slots; + int* modes = owners + slots; + int owner_counts[MEGAMOE_TMA_GPU_PLAN_MAX_WORLD] = {0}; + unsigned char seen_experts[384] = {0}; + int incoming_slots = 0; + int active = 0, common_level = -1; + int global_max_send = plan[1]; + int outgoing_count = plan[3]; + *plan_error = 1; + if (c->plan_abi_version != 6 || c->plan_words != 4 + 7 * stride || plan[2] != 0 || global_max_send < 0 + || global_max_send > quota || outgoing_count < 0 || outgoing_count > quota + || (!c->route_features && outgoing_count)) + return 0; + + /* Published modes are global decisions. Validate the complete local row + * before constructing any descriptor; inactive capacity never becomes work. */ + for (int slot = 0; slot < slots; ++slot) + { + int expert = experts[slot] = plan[4 + slot]; + int level = levels[slot] = plan[4 + stride + slot]; + int owner = owners[slot] = plan[4 + 2 * stride + slot]; + int mode = modes[slot] = plan[4 + 3 * stride + slot]; + if (expert == -1) + { + if (level != -1 || owner != -1 || mode != 0) + return 0; + continue; + } + if (expert < 0 || expert >= c->global_experts || level < 0 || level >= c->level_count || owner < 0 + || owner >= c->world || owner != expert / c->home_count || mode < 0 || mode > 2 || seen_experts[expert]) + return 0; + seen_experts[expert] = 1; + if (++owner_counts[owner] > global_max_send) + return 0; + int group = c->group_sizes[level]; + int begin = c->rank / group * group; + int local = owner >= begin && owner < begin + group; + if ((mode == 1 && !local) + || (mode == 2 && (!c->route_features || !c->foreign_dst_table || local || level == 0))) + return 0; + if (!active) + common_level = level; + else if (common_level != level) + common_level = -2; + ++active; + incoming_slots += mode == 2; + } + /* Source-direct is valid whenever its owner belongs to the target team. + * The scheduler may explicitly choose it for multiple weights or levels; + * auto-policy heuristics are not descriptor safety constraints. */ + if (incoming_slots) + { + if (common_level <= 0 || global_max_send == 0 || global_max_send > 2 || (global_max_send == 2 && active < 3)) + return 0; + int group = c->group_sizes[common_level]; + int begin = c->rank / group * group; + for (int slot = 0; slot < slots; ++slot) + { + if (experts[slot] < 0) + continue; + int external = owners[slot] < begin || owners[slot] >= begin + group; + if (external != (modes[slot] == 2)) + return 0; + } + } + if (outgoing_count && (global_max_send == 0 || global_max_send > 2 || (active && common_level <= 0))) + return 0; + if (owner_counts[c->rank] + outgoing_count > global_max_send) + return 0; + for (int i = 0; i < outgoing_count; ++i) + { + int expert = plan[4 + 4 * stride + i]; + int target = plan[4 + 5 * stride + i]; + int helper = plan[4 + 6 * stride + i]; + int level, begin; + if (expert < c->rank * c->home_count || expert >= (c->rank + 1) * c->home_count || seen_experts[expert] + || helper < 0 || helper >= slots || target_geometry(c, target, &level, &begin)) + return 0; + int group = c->group_sizes[level]; + if (c->rank >= begin && c->rank < begin + group) + return 0; + seen_experts[expert] = 1; + for (int j = 0; j < i; ++j) + if (plan[4 + 5 * stride + j] == target && plan[4 + 6 * stride + j] == helper) + return 0; + } + for (int i = outgoing_count; i < stride; ++i) + if (plan[4 + 4 * stride + i] != -1 || plan[4 + 5 * stride + i] != -1 || plan[4 + 6 * stride + i] != -1) + return 0; + for (int slot = slots; slot < stride; ++slot) + if (plan[4 + slot] != -1 || plan[4 + stride + slot] != -1 || plan[4 + 2 * stride + slot] != -1 + || plan[4 + 3 * stride + slot] != 0) + return 0; + + return build_tma_descriptors(c, outgoing_count, plan_error); +} + +} // namespace megamoe_gpu_plan_detail + +/* Host/device parity permits differential tests against the CPU submitter. + * Counts remain zero on failure: callers must not publish successful READY. + * Caller supplies valid memory for the full immutable tables and scratch. + */ +static MEGAMOE_GP_HD inline int megamoe_tma_build_gpu_plan(MegamoeTmaGpuPlanConfig const* config, + int32_t const* device_plan, int32_t* slot_scratch, MegamoeTmaCopySegment* segments, uint64_t max_segments, + MegamoeTmaCopyRange* ranges, uint64_t max_ranges, MegamoeTmaGpuPlanResult* result) +{ + if (!result) + return 1; + *result = MegamoeTmaGpuPlanResult{}; + result->error = 1; + if (!config || !device_plan || !slot_scratch || !segments || !ranges || !max_segments || !max_ranges + || max_segments > SIZE_MAX || max_ranges > SIZE_MAX) + return 1; + MegamoeTmaGpuPlanConfig const& g = *config; + if (g.world < 2 || g.world > MEGAMOE_TMA_GPU_PLAN_MAX_WORLD || (g.world & 1) || g.rank < 0 || g.rank >= g.world + || g.helper_count <= 0 || g.planes <= 0 || g.global_experts <= 0 || g.global_experts > 384 + || g.global_experts % g.world || g.home_count != g.global_experts / g.world || g.level_count <= 0 + || g.level_count > MEGAMOE_TMA_GPU_PLAN_MAX_LEVELS + || g.owner_stride != ((int64_t(g.helper_count) + 3) & ~INT64_C(3)) + || int64_t(g.plan_words) != 4 + 7 * int64_t(g.owner_stride) || g.plan_abi_version != 6 + || (g.route_features != 0 && g.route_features != 1) || g.tma_route < 0 || g.tma_route > 2 + || g.tma_source_load_percent < 1 || g.tma_source_load_percent > 199 || !g.src_table || !g.dst_table + || !g.plane_bytes || (g.route_features && !g.foreign_dst_table) || (!g.route_features && g.foreign_dst_table) + || device_plan[0] != 0) + return 1; + int targets = 0; + for (int level = 0; level < g.level_count; ++level) + { + int group = g.group_sizes[level]; + if (group < (level ? 4 : 2) || group > g.world || g.world % group || (!level && group != g.world) + || (level && (group >= g.group_sizes[level - 1] || g.group_sizes[level - 1] % group))) + return 1; + if (level) + targets += g.world / group; + } + if (g.target_count != targets || (g.route_features && !targets)) + return 1; + megamoe_gpu_plan_detail::Builder c{}; + static_cast(c) = g; + c.plan = device_plan; + c.slot_scratch = slot_scratch; + c.tma_segments = segments; + c.tma_ranges = ranges; + c.max_commands = (size_t) max_segments; + c.tma_max_ranges = (size_t) max_ranges; + int error = 1; + megamoe_gpu_plan_detail::build_descriptors(&c, &error); + result->error = error; + if (error) + return error; + result->direct_slots = c.last_direct_slots; + result->scatter_slots = c.last_scatter_slots; + result->weighted_domains = c.last_tma_weighted_domains; + result->uniform_domains = c.last_tma_uniform_domains; + result->mixed_source_domains = c.last_tma_mixed_source_domains; + result->foreign_source_domains = c.last_tma_foreign_source_domains; + result->segment_count = c.tma_segment_count; + result->range_count = c.tma_range_count; + result->total_slices = c.last_tma_local_slices; + result->weighted_total_slices = c.last_tma_weighted_total_slices; + result->weighted_source_quota_slices = c.last_tma_weighted_source_quota_slices; + result->weighted_local_slices = c.last_tma_weighted_local_slices; + return 0; +} + +#undef MEGAMOE_GP_HD +#endif diff --git a/cpp/tensorrt_llm/thop/fp4Quantize.cpp b/cpp/tensorrt_llm/thop/fp4Quantize.cpp index 859f1489b2b6..266a825329bf 100644 --- a/cpp/tensorrt_llm/thop/fp4Quantize.cpp +++ b/cpp/tensorrt_llm/thop/fp4Quantize.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2020-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -42,8 +42,9 @@ namespace torch_ext // returns self_fp4, self_block_scale_factors // self_fp4: [M, K / 2], FLOAT4_E2M1X2 // self_block_scale_factors: ceil(M / 128) * 128 * ceil(K / sfVecSize / 4) * 4, SF_DTYPE (UE4M3 or UE8M0) -std::tuple fp4_quantize(at::Tensor const& self, std::optional const& globalScale, - int64_t sfVecSize, bool sfUseUE8M0, bool isSfSwizzledLayout) +std::tuple fp4_quantize_sm_budget(at::Tensor const& self, + std::optional const& globalScale, int64_t sfVecSize, bool sfUseUE8M0, bool isSfSwizzledLayout, + int64_t reservedSms) { CHECK_TH_CUDA(self); CHECK_CONTIGUOUS(self); @@ -88,6 +89,9 @@ std::tuple fp4_quantize(at::Tensor const& self, std::opt = at::detail::empty_cuda({SFSize}, SF_DTYPE, self.device(), /* stride */ std::nullopt); // 1D tensor const thread_local int mMultiProcessorCount = tensorrt_llm::common::getMultiProcessorCount(); + TORCH_CHECK(reservedSms >= 0 && reservedSms < mMultiProcessorCount, + "reserved_sms must be nonnegative and smaller than the device SM count"); + int const availableSms = mMultiProcessorCount - static_cast(reservedSms); auto const layout = isSfSwizzledLayout ? tensorrt_llm::QuantizationSFLayout::SWIZZLED : tensorrt_llm::QuantizationSFLayout::LINEAR; @@ -95,7 +99,7 @@ std::tuple fp4_quantize(at::Tensor const& self, std::opt #define LAUNCH_FP4_QUANTIZE_KERNEL(T, SF_VEC_SIZE) \ tensorrt_llm::kernels::invokeFP4Quantization(1, m, k, reinterpret_cast(self.data_ptr()), \ globalScalePtr, reinterpret_cast(valueE2M1.data_ptr()), \ - reinterpret_cast(scaleFP8SF.data_ptr()), sfUseUE8M0, layout, mMultiProcessorCount, \ + reinterpret_cast(scaleFP8SF.data_ptr()), sfUseUE8M0, layout, availableSms, \ at::cuda::getCurrentCUDAStream(self.get_device())); if (sfUseUE8M0) @@ -158,6 +162,12 @@ std::tuple fp4_quantize(at::Tensor const& self, std::opt return {valueE2M1, scaleFP8SF}; } +std::tuple fp4_quantize(at::Tensor const& self, std::optional const& globalScale, + int64_t sfVecSize, bool sfUseUE8M0, bool isSfSwizzledLayout) +{ + return fp4_quantize_sm_budget(self, globalScale, sfVecSize, sfUseUE8M0, isSfSwizzledLayout, 0); +} + at::Tensor calculate_nvfp4_global_scale(at::Tensor const& input, std::optional const& tokensPerBatch) { CHECK_TH_CUDA(input); @@ -377,6 +387,9 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) m.def( "fp4_quantize(Tensor input, Tensor? globalScale, int sfVecSize, bool sfUseUE8M0=False, bool " "isSfSwizzledLayout=True) -> (Tensor, Tensor)"); + m.def( + "fp4_quantize.sm_budget(Tensor input, Tensor? globalScale, int sfVecSize, bool sfUseUE8M0=False, bool " + "isSfSwizzledLayout=True, int reserved_sms=0) -> (Tensor, Tensor)"); m.def("calculate_nvfp4_global_scale(Tensor input, Tensor? tokensPerBatch) -> Tensor"); m.def( "fp4_quantize_with_reorder_residual(Tensor X, Tensor input_scale, Tensor reorder_index, int KE, bool is_act) " @@ -387,6 +400,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) TORCH_LIBRARY_IMPL(trtllm, CUDA, m) { m.impl("fp4_quantize", TORCH_FN(tensorrt_llm::torch_ext::fp4_quantize)); + m.impl("fp4_quantize.sm_budget", TORCH_FN(tensorrt_llm::torch_ext::fp4_quantize_sm_budget)); m.impl("calculate_nvfp4_global_scale", TORCH_FN(tensorrt_llm::torch_ext::calculate_nvfp4_global_scale)); m.impl("fp4_quantize_with_reorder_residual", TORCH_FN(tensorrt_llm::torch_ext::fp4_quantize_with_reorder_residual)); m.impl("fp4_quantize_with_residual", TORCH_FN(tensorrt_llm::torch_ext::fp4_quantize_with_residual)); diff --git a/cpp/tensorrt_llm/thop/fp8Quantize.cpp b/cpp/tensorrt_llm/thop/fp8Quantize.cpp index e80184ff215c..86388ab8ea10 100644 --- a/cpp/tensorrt_llm/thop/fp8Quantize.cpp +++ b/cpp/tensorrt_llm/thop/fp8Quantize.cpp @@ -151,7 +151,8 @@ std::tuple fp8_batched_quantize_1x128_permute102(at::Ten // so each UE8M0 scale is replicated into four adjacent K32 slots. The legacy // deep_gemm MN-major packed layout remains available while that consumer still // requires it. -std::tuple fp8_quantize_1x128_packed_ue8m0(at::Tensor const& self, bool useR128c4Layout) +std::tuple fp8_quantize_1x128_packed_ue8m0_sm_budget( + at::Tensor const& self, bool useR128c4Layout, int64_t reservedSms) { CHECK_TH_CUDA(self); CHECK_CONTIGUOUS(self); @@ -168,6 +169,8 @@ std::tuple fp8_quantize_1x128_packed_ue8m0(at::Tensor co TORCH_CHECK(n <= std::numeric_limits::max(), "N must be within int32"); auto const num_n_blocks = (n + 127) / 128; auto const num_packed_sf_k = (num_n_blocks + 3) / 4; + TORCH_CHECK(reservedSms >= 0, "reserved_sms must be nonnegative"); + TORCH_CHECK(reservedSms == 0 || useR128c4Layout, "reserved_sms requires the R128c4 scale layout"); auto stream = at::cuda::getCurrentCUDAStream(self.get_device()); if (useR128c4Layout) @@ -188,11 +191,13 @@ std::tuple fp8_quantize_1x128_packed_ue8m0(at::Tensor co { #ifdef ENABLE_BF16 const thread_local int multiProcessorCount = tensorrt_llm::common::getMultiProcessorCount(); + TORCH_CHECK(reservedSms < multiProcessorCount, "reserved_sms must be smaller than the device SM count"); + int const availableSms = multiProcessorCount - static_cast(reservedSms); tensorrt_llm::kernels::invokeMxFP8Quantization<__nv_bfloat16, kQuantSfVecSize, kOutputSfVecSize>(1, static_cast(m), static_cast(n), static_cast(n), reinterpret_cast<__nv_bfloat16 const*>(self.data_ptr()), reinterpret_cast(valueE4M3.data_ptr()), reinterpret_cast(scaleFP8SF.data_ptr()), - tensorrt_llm::QuantizationSFLayout::SWIZZLED, multiProcessorCount, stream); + tensorrt_llm::QuantizationSFLayout::SWIZZLED, availableSms, stream); #else C10_THROW_ERROR(NotImplementedError, "BFloat16 must be enabled to quantize a BF16 tensor to MXFP8."); #endif @@ -232,6 +237,11 @@ std::tuple fp8_quantize_1x128_packed_ue8m0(at::Tensor co return {valueE4M3.slice(0, 0, m), packedScale}; } +std::tuple fp8_quantize_1x128_packed_ue8m0(at::Tensor const& self, bool useR128c4Layout) +{ + return fp8_quantize_1x128_packed_ue8m0_sm_budget(self, useR128c4Layout, 0); +} + std::tuple fp8_quantize_1x128_cutedsl_ue8m0(at::Tensor const& self) { CHECK_TH_CUDA(self); @@ -336,6 +346,9 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) m.def("fp8_quantize_1x128(Tensor input, bool use_ue8m0=False) -> (Tensor, Tensor)"); m.def("fp8_batched_quantize_1x128_permute102(Tensor input) -> (Tensor, Tensor)"); m.def("fp8_quantize_1x128_packed_ue8m0(Tensor input, bool use_r128c4_layout=True) -> (Tensor, Tensor)"); + m.def( + "fp8_quantize_1x128_packed_ue8m0.sm_budget(Tensor input, bool use_r128c4_layout=True, " + "int reserved_sms=0) -> (Tensor, Tensor)"); m.def("fp8_quantize_1x128_cutedsl_ue8m0(Tensor input) -> (Tensor, Tensor)"); m.def( "silu_and_mul_fp8_quantize_1x128_packed_ue8m0(Tensor input, float? swiglu_limit=None, bool " @@ -347,6 +360,8 @@ TORCH_LIBRARY_IMPL(trtllm, CUDA, m) m.impl("fp8_quantize_1x128", &tensorrt_llm::torch_ext::fp8_quantize_1x128); m.impl("fp8_batched_quantize_1x128_permute102", &tensorrt_llm::torch_ext::fp8_batched_quantize_1x128_permute102); m.impl("fp8_quantize_1x128_packed_ue8m0", &tensorrt_llm::torch_ext::fp8_quantize_1x128_packed_ue8m0); + m.impl("fp8_quantize_1x128_packed_ue8m0.sm_budget", + &tensorrt_llm::torch_ext::fp8_quantize_1x128_packed_ue8m0_sm_budget); m.impl("fp8_quantize_1x128_cutedsl_ue8m0", &tensorrt_llm::torch_ext::fp8_quantize_1x128_cutedsl_ue8m0); m.impl("silu_and_mul_fp8_quantize_1x128_packed_ue8m0", &tensorrt_llm::torch_ext::silu_and_mul_fp8_quantize_1x128_packed_ue8m0); diff --git a/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.cpp b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.cpp new file mode 100644 index 000000000000..c1b82109d806 --- /dev/null +++ b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.cpp @@ -0,0 +1,162 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h" + +#include +#include +#include + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +void checkCudaTensor(torch::Tensor const& tensor, c10::Device const& device, c10::ScalarType dtype, + std::int64_t minElements, char const* name) +{ + TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(tensor.device() == device, name, " must be on the routes CUDA device"); + TORCH_CHECK(tensor.scalar_type() == dtype, name, " has an invalid dtype"); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); + TORCH_CHECK(tensor.numel() >= minElements, name, " is too small: expected at least ", minElements, + " elements, got ", tensor.numel()); +} + +int checkedInt(std::int64_t value, char const* name) +{ + TORCH_CHECK(value >= INT_MIN && value <= INT_MAX, name, " does not fit in int32"); + return static_cast(value); +} + +} // namespace + +void moeRebalanceHaloQ(torch::Tensor const& routes, torch::Tensor& outSlots, torch::Tensor& outIds, + torch::Tensor& outLevels, torch::Tensor& outOwners, torch::Tensor const& peerBases, torch::Tensor& status, + torch::Tensor& partial, torch::Tensor& routeAux, torch::Tensor& gridSync, torch::Tensor& planWorkspace, + torch::Tensor& routePrefix, std::int64_t planChannelPtr, std::int64_t epValue, std::int64_t expertsValue, + std::int64_t helpersValue, std::int64_t localRankValue, std::int64_t routeCapacityValue, std::int64_t ctasValue, + std::int64_t threadsValue, bool enablePdl, std::int64_t spinCycles, std::int64_t planAbiVersionValue, + std::int64_t planChannelWordsValue, std::int64_t routeFeaturesValue, std::int64_t validRouteCountValue) +{ + int const ep = checkedInt(epValue, "ep"); + int const experts = checkedInt(expertsValue, "experts"); + int const helpers = checkedInt(helpersValue, "helpers"); + int const localRank = checkedInt(localRankValue, "local_rank"); + int const routeCapacity = checkedInt(routeCapacityValue, "route_capacity"); + int const ctas = checkedInt(ctasValue, "ctas"); + int const threads = checkedInt(threadsValue, "threads"); + int const planAbiVersion = checkedInt(planAbiVersionValue, "plan_abi_version"); + int const planChannelWords = checkedInt(planChannelWordsValue, "plan_channel_words"); + int const routeFeatures = checkedInt(routeFeaturesValue, "route_features"); + int const validRouteCount = checkedInt(validRouteCountValue, "valid_route_count"); + TORCH_CHECK(ep >= 2 && ep <= 32 && ep % 2 == 0, "ep must be even and in [2, 32]"); + TORCH_CHECK(experts > 0 && experts <= 384 && experts % ep == 0, + "experts must be positive, no greater than 384, and divisible by ep"); + TORCH_CHECK(helpers > 0 && helpers <= INT_MAX / ep - experts / ep, "invalid helpers"); + TORCH_CHECK(localRank >= 0 && localRank < ep, "local_rank must be in [0, ep)"); + TORCH_CHECK(routeCapacity > 0 && routeCapacity <= INT_MAX / ep, "invalid route_capacity"); + TORCH_CHECK( + validRouteCount >= 0 && validRouteCount <= routeCapacity, "valid_route_count must be in [0, route_capacity]"); + TORCH_CHECK(ctas > 0 && ctas <= 128 && routeCapacity <= INT_MAX - ctas, "invalid ctas"); + TORCH_CHECK(threads == 512, "threads must be 512"); + TORCH_CHECK(spinCycles >= 0, "spin_cycles cannot be negative"); + TORCH_CHECK(planChannelPtr >= 0, "plan_channel_ptr cannot be negative"); + + c10::Device const device = routes.device(); + checkCudaTensor(routes, device, torch::kInt32, validRouteCount, "routes"); + checkCudaTensor(outSlots, device, torch::kInt32, routeCapacity, "out_slots"); + checkCudaTensor(outIds, device, torch::kInt32, helpers, "out_ids"); + checkCudaTensor(outLevels, device, torch::kInt32, helpers, "out_levels"); + checkCudaTensor(outOwners, device, torch::kInt32, helpers, "out_owners"); + checkCudaTensor(peerBases, device, torch::kInt64, ep, "peer_bases"); + checkCudaTensor(status, device, torch::kInt32, 6, "status"); + checkCudaTensor(partial, device, torch::kInt32, static_cast(ctas) * experts, "partial"); + std::int64_t const bins = static_cast(experts) + 1; + std::int64_t routeAuxElements = static_cast(ctas) * (threads / 32) * bins; + if (ctas > 1) + { + routeAuxElements = std::max(routeAuxElements, 2 + static_cast(ctas - 1) * bins + routeCapacity); + } + checkCudaTensor(routeAux, device, torch::kInt32, routeAuxElements, "route_aux"); + checkCudaTensor(gridSync, device, torch::kInt32, 4, "grid_sync"); + checkCudaTensor(planWorkspace, device, torch::kInt32, 2 + 6 * 384, "plan_workspace"); + checkCudaTensor(routePrefix, device, torch::kInt32, static_cast(384) * (ep + 1), "route_prefix"); + + c10::cuda::CUDAGuard const deviceGuard(device); + auto const stream = at::cuda::getCurrentCUDAStream(routes.get_device()); + kernels::MoeRebalanceHaloQParams const params{ + routes.data_ptr(), + outSlots.data_ptr(), + outIds.data_ptr(), + outLevels.data_ptr(), + outOwners.data_ptr(), + reinterpret_cast(peerBases.data_ptr()), + status.data_ptr(), + partial.data_ptr(), + routeAux.data_ptr(), + gridSync.data_ptr(), + planWorkspace.data_ptr(), + routePrefix.data_ptr(), + reinterpret_cast(static_cast(planChannelPtr)), + ep, + experts, + helpers, + localRank, + routeCapacity, + ctas, + threads, + enablePdl, + static_cast(spinCycles), + planAbiVersion, + planChannelWords, + routeFeatures, + validRouteCount, + }; + kernels::invokeMoeRebalanceHaloQ(params, stream); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "moe_rebalance_halo_q(Tensor routes, Tensor(a!) out_slots, Tensor(b!) " + "out_ids, Tensor(c!) out_levels, " + "Tensor(d!) out_owners, Tensor peer_bases, Tensor(e!) status, " + "Tensor(f!) partial, Tensor(g!) route_aux, " + "Tensor(h!) grid_sync, Tensor(i!) plan_workspace, Tensor(j!) " + "route_prefix, int plan_channel_ptr, int ep, " + "int experts, int helpers, int local_rank, int route_capacity, int " + "ctas, int threads, bool enable_pdl, " + "int spin_cycles, int plan_abi_version, int plan_channel_words, int " + "route_features, " + "int valid_route_count) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("moe_rebalance_halo_q", &tensorrt_llm::torch_ext::moeRebalanceHaloQ); +} diff --git a/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp new file mode 100644 index 000000000000..d88da5d56031 --- /dev/null +++ b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp @@ -0,0 +1,454 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ +namespace +{ + +void checkTmaResult(int error, char const* operation) +{ + TORCH_CHECK(error == 0, operation, " failed: ", megamoe_tma_copy_error_string(error), " (", error, ")"); +} + +int checkedInt(int64_t value, char const* name) +{ + TORCH_CHECK(value >= std::numeric_limits::min() && value <= std::numeric_limits::max(), name, + " is outside the int32 range"); + return static_cast(value); +} + +uint64_t checkedUnsigned(int64_t value, char const* name) +{ + TORCH_CHECK(value >= 0, name, " must be nonnegative"); + return static_cast(value); +} + +std::vector checkedU64Vector(std::vector const& values, uint64_t expectedSize, char const* name) +{ + TORCH_CHECK(expectedSize <= std::numeric_limits::max(), name, " size overflows size_t"); + TORCH_CHECK(values.size() == static_cast(expectedSize), name, " has ", values.size(), " entries; expected ", + expectedSize); + std::vector result; + result.reserve(values.size()); + for (int64_t value : values) + { + result.push_back(checkedUnsigned(value, name)); + } + return result; +} + +void checkBoundTensor(torch::Tensor const& tensor, int device, int64_t minimumElements, char const* name) +{ + TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(tensor.scalar_type() == at::ScalarType::Int, name, " must have dtype torch.int32"); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); + TORCH_CHECK(tensor.get_device() == device, name, " must be on CUDA device ", device); + TORCH_CHECK(tensor.numel() >= minimumElements, name, " has ", tensor.numel(), " elements; expected at least ", + minimumElements); +} + +} // namespace + +class MoeRebalanceTmaState : public torch::CustomClassHolder +{ +public: + MoeRebalanceTmaState(int64_t maxSegments, int64_t sms, int64_t warps) + { + TORCH_CHECK(maxSegments > 0, "max_segments must be positive"); + auto const maxSegmentsUnsigned = checkedUnsigned(maxSegments, "max_segments"); + int const error + = megamoe_tma_copy_create(maxSegmentsUnsigned, checkedInt(sms, "sms"), checkedInt(warps, "warps"), &mState); + checkTmaResult(error, "moe rebalance TMA create"); + + MegamoeTmaCopyConfig config{}; + int const configError = megamoe_tma_copy_config_info(mState, &config); + if (configError != 0) + { + megamoe_tma_copy_destroy(&mState); + checkTmaResult(configError, "moe rebalance TMA config query"); + } + mDevice = config.device; + } + + MoeRebalanceTmaState(MoeRebalanceTmaState const&) = delete; + MoeRebalanceTmaState& operator=(MoeRebalanceTmaState const&) = delete; + + ~MoeRebalanceTmaState() noexcept override + { + std::lock_guard lock(mMutex); + if (mState != nullptr) + { + // Destructors cannot report CUDA cleanup failures. Explicit destroy() is + // available when the caller needs a synchronous error result. + static_cast(megamoe_tma_copy_destroy(&mState)); + } + } + + void configureGpuPlan(int64_t world, int64_t rank, int64_t helperCount, int64_t planes, int64_t globalExperts, + int64_t homeCount, int64_t ownerStride, int64_t levelCount, int64_t targetCount, int64_t planAbiVersion, + int64_t planWords, int64_t routeFeatures, int64_t tmaRoute, int64_t tmaSourceLoadPercent, + std::vector const& groupSizes, std::vector const& sourceTable, + std::vector const& destinationTable, std::vector const& foreignDestinationTable, + std::vector const& planeBytes) + { + std::lock_guard lock(mMutex); + requireLive(); + TORCH_CHECK(!mConfigured, "moe rebalance TMA state is already configured"); + + int const worldValue = checkedInt(world, "world"); + int const helperCountValue = checkedInt(helperCount, "helper_count"); + int const planesValue = checkedInt(planes, "planes"); + int const globalExpertsValue = checkedInt(globalExperts, "global_experts"); + int const levelCountValue = checkedInt(levelCount, "level_count"); + int const targetCountValue = checkedInt(targetCount, "target_count"); + int const routeFeaturesValue = checkedInt(routeFeatures, "route_features"); + TORCH_CHECK(worldValue > 0 && helperCountValue > 0 && planesValue > 0 && globalExpertsValue > 0, + "world, helper_count, planes, and global_experts must be positive"); + TORCH_CHECK(worldValue <= MEGAMOE_TMA_GPU_PLAN_MAX_WORLD && planesValue <= 32 && globalExpertsValue <= 384, + "world, planes, or global_experts exceeds the GPU plan ABI"); + TORCH_CHECK( + helperCountValue <= (std::numeric_limits::max() - 25) / 7, "helper_count exceeds the GPU plan ABI"); + TORCH_CHECK(levelCountValue > 0 && targetCountValue >= 0, + "level_count must be positive and target_count must be nonnegative"); + TORCH_CHECK(targetCountValue <= 32 && (routeFeaturesValue == 0 || routeFeaturesValue == 1), + "target_count or route_features exceeds the GPU plan ABI"); + TORCH_CHECK(levelCountValue <= MEGAMOE_TMA_GPU_PLAN_MAX_LEVELS, "level_count exceeds the GPU plan ABI"); + TORCH_CHECK( + groupSizes.size() == static_cast(levelCountValue), "group_sizes must contain level_count entries"); + + std::array groups{}; + for (int index = 0; index < levelCountValue; ++index) + { + groups[index] = checkedInt(groupSizes[index], "group_sizes entry"); + } + + uint64_t const sourceCount = static_cast(globalExpertsValue) * planesValue; + uint64_t const destinationCount = static_cast(levelCountValue) * helperCountValue * planesValue; + uint64_t const foreignCount + = routeFeaturesValue ? static_cast(targetCountValue) * helperCountValue * planesValue : 0; + auto sources = checkedU64Vector(sourceTable, sourceCount, "source_table"); + auto destinations = checkedU64Vector(destinationTable, destinationCount, "destination_table"); + auto foreign = checkedU64Vector(foreignDestinationTable, foreignCount, "foreign_destination_table"); + auto bytes = checkedU64Vector(planeBytes, planesValue, "plane_bytes"); + + MegamoeTmaGpuPlanConfig config{}; + config.world = worldValue; + config.rank = checkedInt(rank, "rank"); + config.helper_count = helperCountValue; + config.planes = planesValue; + config.global_experts = globalExpertsValue; + config.home_count = checkedInt(homeCount, "home_count"); + config.owner_stride = checkedInt(ownerStride, "owner_stride"); + config.level_count = levelCountValue; + config.target_count = targetCountValue; + config.plan_abi_version = checkedInt(planAbiVersion, "plan_abi_version"); + config.plan_words = checkedInt(planWords, "plan_words"); + config.route_features = routeFeaturesValue; + config.tma_route = checkedInt(tmaRoute, "tma_route"); + config.tma_source_load_percent = checkedInt(tmaSourceLoadPercent, "tma_source_load_percent"); + for (size_t index = 0; index < groups.size(); ++index) + { + config.group_sizes[index] = groups[index]; + } + config.src_table = sources.data(); + config.dst_table = destinations.data(); + config.foreign_dst_table = foreign.empty() ? nullptr : foreign.data(); + config.plane_bytes = bytes.data(); + + c10::cuda::CUDAGuard const deviceGuard(c10::Device(c10::DeviceType::CUDA, mDevice)); + checkTmaResult(megamoe_tma_copy_configure_gpu_plan(mState, &config), "moe rebalance TMA GPU plan configure"); + mConfigured = true; + mHelperCount = helperCountValue; + } + + void bindGpuDirect(torch::Tensor ids, torch::Tensor levels, torch::Tensor owners, torch::Tensor workspace, + torch::Tensor status, torch::Tensor gridSync, int64_t capacity) + { + std::lock_guard lock(mMutex); + requireLive(); + TORCH_CHECK(mConfigured, "moe rebalance TMA state must be configured before bind"); + TORCH_CHECK(!mBound, "moe rebalance TMA state is already bound"); + int const capacityValue = checkedInt(capacity, "capacity"); + TORCH_CHECK(capacityValue > 0, "capacity must be positive"); + checkBoundTensor(ids, mDevice, mHelperCount, "ids"); + checkBoundTensor(levels, mDevice, mHelperCount, "levels"); + checkBoundTensor(owners, mDevice, mHelperCount, "owners"); + checkBoundTensor(workspace, mDevice, 2 + 6LL * capacityValue, "workspace"); + checkBoundTensor(status, mDevice, 6, "status"); + checkBoundTensor(gridSync, mDevice, 2, "grid_sync"); + + c10::cuda::CUDAGuard const deviceGuard(ids.device()); + checkTmaResult( + megamoe_tma_copy_bind_gpu_direct(mState, reinterpret_cast(ids.data_ptr()), + reinterpret_cast(levels.data_ptr()), reinterpret_cast(owners.data_ptr()), + reinterpret_cast(workspace.data_ptr()), reinterpret_cast(status.data_ptr()), + reinterpret_cast(gridSync.data_ptr()), capacityValue), + "moe rebalance TMA GPU-direct bind"); + + // The native state stores raw device addresses. Retaining the tensors in + // the owner prevents allocator reuse until destroy() has drained work. + mIds = std::move(ids); + mLevels = std::move(levels); + mOwners = std::move(owners); + mWorkspace = std::move(workspace); + mStatus = std::move(status); + mGridSync = std::move(gridSync); + mBound = true; + } + + int64_t submitGpuDirect(int64_t flagMc) + { + std::lock_guard lock(mMutex); + requireBound(); + c10::cuda::CUDAGuard const deviceGuard(c10::Device(c10::DeviceType::CUDA, mDevice)); + auto const stream = at::cuda::getCurrentCUDAStream(mDevice).stream(); + submitLocked(flagMc, reinterpret_cast(stream)); + return 0; + } + + int64_t currentGeneration() const + { + std::lock_guard lock(mMutex); + requireLive(); + return checkedDiagnostic(mGeneration); + } + + // Result order: error, direct/scatter slots, four domain counts, segment/range/slice counts, + // weighted total/quota/local slices. This diagnostic synchronizes the last submitted stream. + + std::vector result() + { + std::lock_guard lock(mMutex); + requireLive(); + TORCH_CHECK(mConfigured, "moe rebalance TMA state is not configured"); + c10::cuda::CUDAGuard const deviceGuard(c10::Device(c10::DeviceType::CUDA, mDevice)); + MegamoeTmaGpuPlanResult value{}; + checkTmaResult(megamoe_tma_copy_gpu_plan_result(mState, &value), "moe rebalance TMA diagnostic result"); + return {value.error, value.direct_slots, value.scatter_slots, value.weighted_domains, value.uniform_domains, + value.mixed_source_domains, value.foreign_source_domains, checkedDiagnostic(value.segment_count), + checkedDiagnostic(value.range_count), checkedDiagnostic(value.total_slices), + checkedDiagnostic(value.weighted_total_slices), checkedDiagnostic(value.weighted_source_quota_slices), + checkedDiagnostic(value.weighted_local_slices)}; + } + + // Config order follows MegamoeTmaCopyConfig through max_segments. + + std::vector config() const + { + std::lock_guard lock(mMutex); + requireLive(); + MegamoeTmaCopyConfig value{}; + checkTmaResult(megamoe_tma_copy_config_info(mState, &value), "moe rebalance TMA config query"); + return {value.abi_version, value.device, value.device_sm_count, value.sms, value.warps, value.threads_per_cta, + value.slots_per_warp, value.bank0_slots_per_warp, value.bank1_slots_per_warp, value.slice_bytes, + value.dynamic_shared_bytes, value.device_optin_shared_bytes, value.device_shared_bytes_per_sm, + value.max_active_ctas_per_sm, value.compute_major, value.compute_minor, value.total_slots, + value.max_slots_per_warp, value.extra_slot_warps, value.max_warps, checkedDiagnostic(value.max_segments)}; + } + + void destroy() + { + std::lock_guard lock(mMutex); + if (mState == nullptr) + { + return; + } + c10::cuda::CUDAGuard const deviceGuard(c10::Device(c10::DeviceType::CUDA, mDevice)); + int const error = megamoe_tma_copy_destroy(&mState); + if (mState == nullptr) + { + mIds = torch::Tensor(); + mLevels = torch::Tensor(); + mOwners = torch::Tensor(); + mWorkspace = torch::Tensor(); + mStatus = torch::Tensor(); + mGridSync = torch::Tensor(); + mConfigured = false; + mBound = false; + } + checkTmaResult(error, "moe rebalance TMA destroy"); + } + +private: + void requireLive() const + { + TORCH_CHECK(mState != nullptr, "moe rebalance TMA state has been destroyed"); + } + + void requireBound() const + { + requireLive(); + TORCH_CHECK(mBound, "moe rebalance TMA state must be bound before submit"); + } + + static int64_t checkedDiagnostic(uint64_t value) + { + TORCH_CHECK(value <= static_cast(std::numeric_limits::max()), + "moe rebalance TMA diagnostic counter exceeds int64"); + return static_cast(value); + } + + void submitLocked(int64_t flagMc, void* stream) + { + uint64_t const flag = checkedUnsigned(flagMc, "flag_mc"); + TORCH_CHECK(flag != 0, "flag_mc must be nonzero"); + checkTmaResult(megamoe_tma_copy_submit_gpu_direct(mState, flag, mGeneration, stream), + "moe rebalance TMA GPU-direct submit"); + ++mGeneration; + TORCH_CHECK(mGeneration != 0, "moe rebalance TMA generation overflow"); + } + + mutable std::mutex mMutex; + MegamoeTmaCopyState* mState{nullptr}; + int mDevice{-1}; + int mHelperCount{0}; + uint64_t mGeneration{1}; + bool mConfigured{false}; + bool mBound{false}; + torch::Tensor mIds; + torch::Tensor mLevels; + torch::Tensor mOwners; + torch::Tensor mWorkspace; + torch::Tensor mStatus; + torch::Tensor mGridSync; +}; + +using MoeRebalanceTmaStatePtr = c10::intrusive_ptr; + +MoeRebalanceTmaStatePtr moeRebalanceTmaCreate(int64_t maxSegments, int64_t sms, int64_t warps) +{ + return c10::make_intrusive(maxSegments, sms, warps); +} + +void moeRebalanceTmaConfigureGpuPlan(MoeRebalanceTmaStatePtr const& state, int64_t world, int64_t rank, + int64_t helperCount, int64_t planes, int64_t globalExperts, int64_t homeCount, int64_t ownerStride, + int64_t levelCount, int64_t targetCount, int64_t planAbiVersion, int64_t planWords, int64_t routeFeatures, + int64_t tmaRoute, int64_t tmaSourceLoadPercent, std::vector const& groupSizes, + std::vector const& sourceTable, std::vector const& destinationTable, + std::vector const& foreignDestinationTable, std::vector const& planeBytes) +{ + TORCH_CHECK(state, "moe rebalance TMA state is null"); + state->configureGpuPlan(world, rank, helperCount, planes, globalExperts, homeCount, ownerStride, levelCount, + targetCount, planAbiVersion, planWords, routeFeatures, tmaRoute, tmaSourceLoadPercent, groupSizes, sourceTable, + destinationTable, foreignDestinationTable, planeBytes); +} + +void moeRebalanceTmaBindGpuDirect(MoeRebalanceTmaStatePtr const& state, torch::Tensor ids, torch::Tensor levels, + torch::Tensor owners, torch::Tensor workspace, torch::Tensor status, torch::Tensor gridSync, int64_t capacity) +{ + TORCH_CHECK(state, "moe rebalance TMA state is null"); + state->bindGpuDirect(std::move(ids), std::move(levels), std::move(owners), std::move(workspace), std::move(status), + std::move(gridSync), capacity); +} + +int64_t moeRebalanceTmaSubmitGpuDirect(MoeRebalanceTmaStatePtr const& state, int64_t flagMc) +{ + TORCH_CHECK(state, "moe rebalance TMA state is null"); + return state->submitGpuDirect(flagMc); +} + +int64_t moeRebalanceTmaCurrentGeneration(MoeRebalanceTmaStatePtr const& state) +{ + TORCH_CHECK(state, "moe rebalance TMA state is null"); + return state->currentGeneration(); +} + +std::vector moeRebalanceTmaResult(MoeRebalanceTmaStatePtr const& state) +{ + TORCH_CHECK(state, "moe rebalance TMA state is null"); + return state->result(); +} + +std::vector moeRebalanceTmaConfig(MoeRebalanceTmaStatePtr const& state) +{ + TORCH_CHECK(state, "moe rebalance TMA state is null"); + return state->config(); +} + +void moeRebalanceTmaDestroy(MoeRebalanceTmaStatePtr const& state) +{ + TORCH_CHECK(state, "moe rebalance TMA state is null"); + state->destroy(); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.class_("MoeRebalanceTmaState"); + m.def( + "moe_rebalance_tma_create(int max_segments, int sms=0, int warps=0) -> " + "__torch__.torch.classes.trtllm.MoeRebalanceTmaState"); + m.def( + "moe_rebalance_tma_configure_gpu_plan(" + "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state, int world, int rank, int helper_count, " + "int planes, int global_experts, int home_count, int owner_stride, int level_count, int target_count, " + "int plan_abi_version, int plan_words, int route_features, int tma_route, int tma_source_load_percent, " + "int[] group_sizes, int[] source_table, int[] destination_table, int[] foreign_destination_table, " + "int[] plane_bytes) -> ()"); + m.def( + "moe_rebalance_tma_bind_gpu_direct(" + "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state, Tensor ids, Tensor levels, Tensor owners, " + "Tensor workspace, Tensor status, Tensor grid_sync, int capacity) -> ()"); + m.def( + "moe_rebalance_tma_submit_gpu_direct(" + "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state, int flag_mc) -> int"); + m.def( + "moe_rebalance_tma_current_gen(" + "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state) -> int"); + m.def( + "moe_rebalance_tma_result(" + "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state) -> int[]"); + m.def( + "moe_rebalance_tma_config(" + "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state) -> int[]"); + m.def( + "moe_rebalance_tma_destroy(" + "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CompositeExplicitAutograd, m) +{ + m.impl("moe_rebalance_tma_create", &tensorrt_llm::torch_ext::moeRebalanceTmaCreate); + m.impl("moe_rebalance_tma_configure_gpu_plan", &tensorrt_llm::torch_ext::moeRebalanceTmaConfigureGpuPlan); + m.impl("moe_rebalance_tma_bind_gpu_direct", &tensorrt_llm::torch_ext::moeRebalanceTmaBindGpuDirect); + m.impl("moe_rebalance_tma_submit_gpu_direct", &tensorrt_llm::torch_ext::moeRebalanceTmaSubmitGpuDirect); + m.impl("moe_rebalance_tma_current_gen", &tensorrt_llm::torch_ext::moeRebalanceTmaCurrentGeneration); + m.impl("moe_rebalance_tma_result", &tensorrt_llm::torch_ext::moeRebalanceTmaResult); + m.impl("moe_rebalance_tma_config", &tensorrt_llm::torch_ext::moeRebalanceTmaConfig); + m.impl("moe_rebalance_tma_destroy", &tensorrt_llm::torch_ext::moeRebalanceTmaDestroy); +} diff --git a/pyproject.toml b/pyproject.toml index bbbe3144ae5f..ea7b9f5d70cb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,6 +49,9 @@ exclude = [ "tensorrt_llm/_torch/moe/kernels/rubin/manual_mma_128dp.py", "tensorrt_llm/_torch/moe/kernels/rubin/rubin_contiguous_grouped_blockscaled_gemm_fused_fc12.py", "tensorrt_llm/grpc/openengine/_generated/**", + # Keep pinned MegaMoE sources byte-identical to upstream. + "tensorrt_llm/_torch/cute_dsl_kernels/cutedsl_megamoe/**", + "tensorrt_llm/_torch/cute_dsl_kernels/megamoe_scheduler/**", ] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_10x__default/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_10x__default/ops/moe/__init__.py index 52a7a9daf028..c72d7918ac70 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_10x__default/ops/moe/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_10x__default/ops/moe/__init__.py @@ -1,2 +1,25 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""NVLink-domain pointer mapping and token communication.""" + +from ...helpers.software_sync import NvlinkBarrier, SoftwareGridSync +from ...quant_def import QuantKind +from .symmetric_buffer import SymmetricBufferDevice, SymmetricBufferHost +from .token_comm import ( + TokenBackMode, + TokenBackScheduleMode, + TokenCommArgs, + TokenCommNonDeterministic, +) +from .token_comm_deterministic import TokenCommDeterministic + +__all__ = [ + "NvlinkBarrier", + "QuantKind", + "SoftwareGridSync", + "SymmetricBufferDevice", + "SymmetricBufferHost", + "TokenBackMode", + "TokenBackScheduleMode", + "TokenCommArgs", + "TokenCommDeterministic", + "TokenCommNonDeterministic", +] diff --git a/tensorrt_llm/_torch/attention/kernels/rubin/dsv4_qb_fusion/__init__.py b/tensorrt_llm/_torch/attention/kernels/rubin/dsv4_qb_fusion/__init__.py index 52a7a9daf028..cc1ef51580d3 100644 --- a/tensorrt_llm/_torch/attention/kernels/rubin/dsv4_qb_fusion/__init__.py +++ b/tensorrt_llm/_torch/attention/kernels/rubin/dsv4_qb_fusion/__init__.py @@ -1,2 +1,22 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Cross-rank communication protocols and implementations.""" + +from ..quant_def import CombineFormat, QuantKind +from .nvlink_domain.token_comm import ( + TokenBackScheduleMode, + TokenBackMode, + TokenCommArgs, + TokenCommNonDeterministic, +) +from .nvlink_domain.token_comm_deterministic import TokenCommDeterministic +from .token_protocol import TokenSrcMetadata + +__all__ = [ + "CombineFormat", + "QuantKind", + "TokenBackScheduleMode", + "TokenBackMode", + "TokenCommArgs", + "TokenCommDeterministic", + "TokenCommNonDeterministic", + "TokenSrcMetadata", +] diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index ad4bcbb77d1c..394c5625d371 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -72,6 +72,37 @@ def copy_kv_slots_( ) -> None: return None + @torch.library.register_fake("trtllm::moe_rebalance_halo_q") + def _( + routes: torch.Tensor, + out_slots: torch.Tensor, + out_ids: torch.Tensor, + out_levels: torch.Tensor, + out_owners: torch.Tensor, + peer_bases: torch.Tensor, + status: torch.Tensor, + partial: torch.Tensor, + route_aux: torch.Tensor, + grid_sync: torch.Tensor, + plan_workspace: torch.Tensor, + route_prefix: torch.Tensor, + plan_channel_ptr: int, + ep: int, + experts: int, + helpers: int, + local_rank: int, + route_capacity: int, + ctas: int, + threads: int, + enable_pdl: bool, + spin_cycles: int, + plan_abi_version: int, + plan_channel_words: int, + route_features: int, + valid_route_count: int, + ) -> None: + return None + @torch.library.register_fake("trtllm::allreduce") def allreduce( input: torch.Tensor, @@ -630,7 +661,7 @@ def _(activation: torch.Tensor): activation.new_empty(scale_shape, dtype=activation.dtype)) @torch.library.register_fake("trtllm::fp4_quantize") - def _( + def _fp4_quantize_fake( input: torch.Tensor, global_scale: torch.Tensor, sf_vec_size: int, @@ -643,6 +674,18 @@ def _( return (input.new_empty(output_shape, dtype=torch.uint8), global_scale.new_empty(scale_shape, dtype=torch.uint8)) + if hasattr(torch.ops.trtllm.fp4_quantize, "sm_budget"): + + @torch.library.register_fake("trtllm::fp4_quantize.sm_budget") + def _(input: torch.Tensor, + global_scale: torch.Tensor, + sf_vec_size: int, + sf_use_ue8m0: bool = False, + swizzled_layout: bool = True, + reserved_sms: int = 0): + return _fp4_quantize_fake(input, global_scale, sf_vec_size, + sf_use_ue8m0, swizzled_layout) + @torch.library.register_fake("trtllm::fp4_quantize_with_reorder_residual") def _( X: torch.Tensor, @@ -1107,8 +1150,10 @@ def _(pe: torch.Tensor, nope: torch.Tensor): return packed, scale @torch.library.register_fake("trtllm::fp8_quantize_1x128_packed_ue8m0") - def _(input: torch.Tensor, - use_r128c4_layout: bool = True) -> tuple[torch.Tensor, torch.Tensor]: + def _fp8_quantize_packed_fake( + input: torch.Tensor, + use_r128c4_layout: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: m, k = input.shape[0], input.shape[1] num_n_blocks = (k + 127) // 128 num_packed_sf_k = (num_n_blocks + 3) // 4 @@ -1123,6 +1168,15 @@ def _(input: torch.Tensor, dtype=torch.int32) return torch.empty_like(input, dtype=torch.float8_e4m3fn), scale + if hasattr(torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0, "sm_budget"): + + @torch.library.register_fake( + "trtllm::fp8_quantize_1x128_packed_ue8m0.sm_budget") + def _(input: torch.Tensor, + use_r128c4_layout: bool = True, + reserved_sms: int = 0): + return _fp8_quantize_packed_fake(input, use_r128c4_layout) + @torch.library.register_fake("trtllm::fp8_quantize_1x128_cutedsl_ue8m0") def _(input: torch.Tensor): m, k = input.shape diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index a29eae1fa19e..dc5ecd8f7bde 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -39,10 +39,13 @@ torch_dtype_to_binding) from tensorrt_llm.bindings import LayerType as LayerTypeCpp from tensorrt_llm.functional import AllReduceStrategy +# Callers import MoeLoadBalancerConfig from this module. +from tensorrt_llm.llmapi.llm_args import MoeLoadBalancerConfig # noqa: F401 from tensorrt_llm.llmapi.llm_args import (CSA2SparseAttentionConfig, DeepSeekSparseAttentionConfig, DeepSeekV4SparseAttentionConfig, - KvCacheConfig, MoeLoadBalancerConfig, + KvCacheConfig, + MoeLoadBalancerConfigType, MultimodalConfig) from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping @@ -286,7 +289,7 @@ class ModelConfig(Generic[TConfig]): _moe_max_num_tokens_is_default: Optional[bool] = field(default=None, repr=False, compare=False) - moe_load_balancer: Optional[MoeLoadBalancerConfig] = None + moe_load_balancer: Optional[MoeLoadBalancerConfigType] = None attn_backend: str = 'TRTLLM' # Effective threshold requested for trtllm-gen MLA skip-correction. Zero diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 9932d64e7014..c516a7bb5cf8 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -72,7 +72,7 @@ from ..modules.gated_mlp import GatedMLP from ..modules.linear import Linear, TensorParallelMode, WeightsLoadingConfig from ..modules.mhc.hyper_connection import HCHead, HCState, mHC -from ..modules.multi_stream_utils import maybe_execute_in_parallel +from ..modules.multi_stream_utils import do_multi_stream, maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm from ..moe.fused_moe import ( DEFAULT_MOE_ACTIVATION, @@ -1697,25 +1697,75 @@ def __init__( self.mapping = model_config.mapping - # FIXME: incompatible with mixed quantization mode (including excluding modules from quantization) - block_size = 1 - if model_config.quant_config and model_config.quant_config.group_size is not None: - block_size = model_config.quant_config.group_size + gate_up_quant_config, down_quant_config = self._get_shared_expert_projection_quant_configs( + model_config, layer_idx + ) + shared_model_config = model_config + if ( + gate_up_quant_config == down_quant_config + and gate_up_quant_config is not model_config.quant_config + ): + shared_model_config = copy.copy(model_config) + shared_model_config.quant_config = gate_up_quant_config + quantized_block_sizes = [ + quant_config.group_size + for quant_config in (gate_up_quant_config, down_quant_config) + if ( + quant_config is not None + and quant_config.quant_algo is not None + and quant_config.group_size is not None + ) + ] + block_size = max(quantized_block_sizes, default=1) shared_tp_size, self.shared_output_scale = self._compute_shared_expert_tp_size( shared_expert_intermediate_size, block_size ) - self.shared_experts = GatedMLP( + load_balancer_config = getattr(model_config, "moe_load_balancer", None) + self._per_layer_dynamic_eplb_enabled = ( + getattr(load_balancer_config, "mode", None) == "per_layer_dynamic" + ) + self._per_layer_dynamic_eplb_auxiliary_sms = ( + int(load_balancer_config.auxiliary_sms) if self._per_layer_dynamic_eplb_enabled else 0 + ) + rebalance_backend = getattr(self.experts, "backend", self.experts) + # The fused shared kernel is part of the opt-in per-layer dynamic path. Ask + # the selected routed backend for the capability instead of naming a + # concrete implementation here. + use_fused_fc12 = ( + self._per_layer_dynamic_eplb_enabled + and self._shared_fc12_backend_supported(rebalance_backend) + and get_sm_version() == 107 + and self._shared_fc12_quantization_supported(gate_up_quant_config, down_quant_config) + and not model_config.use_cuda_graph + ) + + shared_mlp_cls = GatedMLP + if use_fused_fc12: + from ..modules.megamoe_shared_mlp import MegaMoESharedMLP + + shared_mlp_cls = MegaMoESharedMLP + + # Only fused FC12 accepts an SM budget; other shared MLP backends stay unrestricted. + shared_reserved_sms = ( + self._per_layer_dynamic_eplb_auxiliary_sms + if use_fused_fc12 and self._per_layer_dynamic_eplb_enabled + else 0 + ) + shared_mlp_kwargs = {"fc12_reserved_sms": shared_reserved_sms} if use_fused_fc12 else {} + + self.shared_experts = shared_mlp_cls( hidden_size=hidden_size, intermediate_size=shared_expert_intermediate_size, bias=False, dtype=dtype, - config=model_config, + config=shared_model_config, overridden_tp_size=shared_tp_size, reduce_output=False, use_cute_dsl_blockscaling_mm=model_config.use_cute_dsl_blockscaling_mm, swiglu_limit=swiglu_limit, + **shared_mlp_kwargs, ) self.allreduce = None @@ -1788,6 +1838,62 @@ def _get_experts_quant_config(model_config, layer_idx: int) -> QuantConfig: return model_config.quant_config return experts_quant_config + @staticmethod + def _get_shared_expert_projection_quant_configs( + model_config, layer_idx: int + ) -> tuple[Optional[QuantConfig], Optional[QuantConfig]]: + base_name = f"model.layers.{layer_idx}.mlp.shared_experts" + quant_config_dict = getattr(model_config, "quant_config_dict", None) or {} + global_quant_config = model_config.quant_config + + gate_up_name = f"{base_name}.gate_up_proj" + gate_name = f"{base_name}.gate_proj" + up_name = f"{base_name}.up_proj" + down_name = f"{base_name}.down_proj" + + # Match apply_layerwise_quant_config exactly for the actual fused + # gate_up Linear: the first gate_proj/gate_up_proj entry wins, while an + # up-only entry does not independently reconfigure a fused weight. + gate_up_quant_config = global_quant_config + for name, candidate in quant_config_dict.items(): + if gate_name in name or gate_up_name in name: + gate_up_quant_config = candidate + break + down_quant_config = quant_config_dict.get(down_name, global_quant_config) + + # apply_quant_config_exclude_modules runs after layerwise overrides and + # excludes the fused gate_up Linear when any fused constituent matches. + if global_quant_config is not None: + unquantized = QuantConfig( + quant_algo=None, + kv_cache_quant_algo=global_quant_config.kv_cache_quant_algo, + ) + if any( + global_quant_config.is_module_excluded_from_quantization(name) + for name in (gate_up_name, gate_name, up_name) + ): + gate_up_quant_config = unquantized + if global_quant_config.is_module_excluded_from_quantization(down_name): + down_quant_config = unquantized + return gate_up_quant_config, down_quant_config + + @staticmethod + def _shared_fc12_backend_supported(backend) -> bool: + capabilities = getattr(backend, "capabilities", None) + return bool(getattr(capabilities, "supports_per_layer_dynamic_eplb", False)) + + @staticmethod + def _shared_fc12_quantization_supported( + gate_up_quant_config: Optional[QuantConfig], + down_quant_config: Optional[QuantConfig], + ) -> bool: + return ( + gate_up_quant_config is not None + and gate_up_quant_config == down_quant_config + and gate_up_quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES + and gate_up_quant_config.group_size == 128 + ) + def compute_routed_output( self, hidden_states, @@ -1874,13 +1980,48 @@ def _compute_routed_output(): # NOTE: define compiled helpers at module scope to avoid defining decorators inside compiled frames - routed_output, shared_output = maybe_execute_in_parallel( - _compute_routed_output, - _compute_shared_output, - self.event_dict[EventType.Main], - self.event_dict[EventType.MoeShared], - self.aux_stream, - ) + # With helper slots enabled, the one-shot hook runs after MAIN has + # enqueued both HALO-Q and TMA copy on the shared high-priority stream. + # Shared experts then run on MAIN while helper weights are copied. + # Install identically across ranks and clear the hook on every exit. + _plan_gap_installed = False + if ( + self._per_layer_dynamic_eplb_enabled + and self.shared_experts is not None + and not do_multi_stream() + ): + _shared_box = {} + + def _plan_gap_shared(): + # Shared-expert scaling must run at most once per forward. + if "out" not in _shared_box: + _shared_box["out"] = _compute_shared_output() + + _experts = getattr(self, "experts", None) + if _experts is not None and hasattr(_experts, "__dict__"): + setattr(_experts, "_rebalance_plan_gap_hook", _plan_gap_shared) + _plan_gap_installed = True + + if _plan_gap_installed: + try: + routed_output = _compute_routed_output() + finally: + # Clear the hook even when the routed path exits early. + _e = getattr(self, "experts", None) + if _e is not None: + setattr(_e, "_rebalance_plan_gap_hook", None) + # Compute shared experts here if the routed path skipped the hook. + if "out" not in _shared_box: + _shared_box["out"] = _compute_shared_output() + shared_output = _shared_box["out"] + else: + routed_output, shared_output = maybe_execute_in_parallel( + _compute_routed_output, + _compute_shared_output, + self.event_dict[EventType.Main], + self.event_dict[EventType.MoeShared], + self.aux_stream, + ) if not do_finalize: return [shared_output, *routed_output] diff --git a/tensorrt_llm/_torch/modules/megamoe_shared_mlp.py b/tensorrt_llm/_torch/modules/megamoe_shared_mlp.py new file mode 100644 index 000000000000..34bef124376c --- /dev/null +++ b/tensorrt_llm/_torch/modules/megamoe_shared_mlp.py @@ -0,0 +1,243 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared FP8 MLP using a single MegaMoE FC1/FC2 compute kernel.""" + +from collections.abc import Callable + +import torch +import torch.nn.functional as F + +from ..distributed import AllReduceParams +from ..model_config import ModelConfig +from .gated_mlp import GatedMLP + + +def _unpack_r128c4(scales: torch.Tensor, rows: int, columns: int) -> torch.Tensor: + padded_rows = (rows + 127) // 128 * 128 + padded_columns = (columns + 3) // 4 * 4 + if scales.dtype != torch.uint8 or scales.ndim != 1: + raise ValueError("Shared FC12 requires loaded uint8 R128c4 weight scales") + if not scales.is_contiguous() or scales.numel() != padded_rows * padded_columns: + raise ValueError("Shared FC12 weight scale storage has an incompatible shape") + return ( + scales.reshape(padded_rows // 128, padded_columns // 4, 32, 4, 4) + .permute(0, 3, 2, 1, 4) + .reshape(padded_rows, padded_columns)[:rows, :columns] + .contiguous() + ) + + +def _pack_r128c4(scales: torch.Tensor) -> torch.Tensor: + rows, columns = scales.shape + padded = F.pad(scales, (0, -columns % 4, 0, -rows % 128)) + return ( + padded.reshape((rows + 127) // 128, 4, 32, (columns + 3) // 4, 4) + .permute(0, 3, 2, 1, 4) + .contiguous() + .reshape(1, -1) + ) + + +def _interleave_gate_up(tensor: torch.Tensor) -> torch.Tensor: + intermediate = tensor.shape[0] // 2 + # Shared Linear stores gate || up; the FC12 epilogue consumes alternating 16-row blocks. + return ( + tensor.reshape(2, intermediate // 16, 16, tensor.shape[1]) + .permute(1, 0, 2, 3) + .contiguous() + .reshape(tensor.shape) + ) + + +class MegaMoESharedMLP(GatedMLP): + """Preserve GatedMLP loading and TP shards while fusing both projections. + + Linear owns checkpoint loading and its existing FP8 resmoothing. This adapter + only rearranges the loaded FP8/UE8M0 bytes. The caller owns shared-output + scaling and reduction; this module always returns the local TP contribution. + """ + + def __init__( + self, + *, + hidden_size: int, + intermediate_size: int, + bias: bool, + activation: Callable[[torch.Tensor], torch.Tensor] = F.silu, + dtype: torch.dtype | None = None, + config: ModelConfig | None = None, + overridden_tp_size: int | None = None, + reduce_output: bool = False, + layer_idx: int | None = None, + use_cute_dsl_blockscaling_mm: bool = False, + disable_deep_gemm: bool = False, + use_custom_cublas_mm: bool = False, + is_shared_expert: bool = False, + swiglu_limit: float | None = None, + swiglu_alpha: float | None = None, + swiglu_beta: float | None = None, + fc12_reserved_sms: int = 0, + ): + if bias or activation is not F.silu: + raise ValueError("Shared FC12 requires bias=False and SwiGLU activation") + if swiglu_alpha not in (None, 1.0) or swiglu_beta not in (None, 0.0): + raise ValueError("Shared FC12 requires SwiGLU alpha=1 and beta=0") + if swiglu_limit is not None and not float(swiglu_limit) > 0: + raise ValueError("Shared FC12 clamp must be positive or None") + if dtype not in (None, torch.bfloat16): + raise ValueError("Shared FC12 requires BF16 activations") + if reduce_output or use_custom_cublas_mm: + raise ValueError("Shared FC12 requires local output and the CuTe DSL backend") + if type(fc12_reserved_sms) is not int or fc12_reserved_sms < 0: + raise ValueError("Shared FC12 reserved-SM budget must be a non-negative integer") + config = config or ModelConfig() + if config.use_cuda_graph: + raise ValueError("Shared FC12 currently requires eager execution") + super().__init__( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + bias=bias, + activation=activation, + dtype=dtype, + config=config, + overridden_tp_size=overridden_tp_size, + reduce_output=False, + layer_idx=layer_idx, + # This selects the existing Linear post-load UE8M0 conversion, including when + # the caller's dense GEMM preference is unset. + use_cute_dsl_blockscaling_mm=True, + disable_deep_gemm=disable_deep_gemm, + is_shared_expert=is_shared_expert, + swiglu_limit=swiglu_limit, + swiglu_alpha=swiglu_alpha, + swiglu_beta=swiglu_beta, + ) + self._validate_fp8_quantization() + self._fc12_reserved_sms = fc12_reserved_sms + self._fc12_max_tokens = int(config.max_num_tokens) + if self._fc12_max_tokens <= 0: + raise ValueError("Shared FC12 max_num_tokens must be positive") + self._fc12_source_state: tuple[tuple[int, int, int | None], ...] | None = None + for name in ("fc1_weight", "fc1_weight_sf", "fc2_weight", "fc2_weight_sf"): + self.register_buffer(f"_fc12_{name}", None, persistent=False) + + backend_name = "cutedsl_megamoe_shared_fc12" + kernel_name = "BlockScaledSwapAbFc12Kernel" + quantization_kind = "mxfp8_e4m3" + + def _validate_fp8_quantization(self) -> None: + for projection in (self.gate_up_proj, self.down_proj): + quant_config = projection.quant_config + if quant_config is None or not quant_config.layer_quant_mode.has_fp8_block_scales(): + raise ValueError("Shared FC12 requires FP8_BLOCK_SCALES for both projections") + + @property + def sm_count(self) -> int: + if not self.gate_up_proj._weights_created: + raise RuntimeError("Shared FC12 SM count requires materialized CUDA weights") + device = self.gate_up_proj.weight.device + if device.type != "cuda": + raise RuntimeError("Shared FC12 SM count requires CUDA weights") + total_sms = torch.cuda.get_device_properties(device).multi_processor_count + reserved = self._fc12_reserved_sms + if not 0 <= reserved < total_sms: + raise ValueError("Shared FC12 reserved-SM budget must leave at least one SM") + return total_sms - reserved + + def cache_derived_state(self) -> None: + """Drop derived views after loading, including staged in-place restoration.""" + self._fc12_source_state = None + for name in ("fc1_weight", "fc1_weight_sf", "fc2_weight", "fc2_weight_sf"): + setattr(self, f"_fc12_{name}", None) + + def post_load_weights(self) -> None: + self.cache_derived_state() + + @torch.no_grad() + def prepare_fc12_weights(self) -> None: + """Cache FC12 layouts after both Linear post-load transformations finish.""" + gate_up, down = self.gate_up_proj, self.down_proj + self._validate_fp8_quantization() + if not gate_up._weights_transformed or not down._weights_transformed: + raise RuntimeError("Shared FC12 requires both Linear post_load_weights calls first") + sources = (gate_up.weight, gate_up.weight_scale, down.weight, down.weight_scale) + source_state = tuple( + (id(t), t.data_ptr(), None if t.is_inference() else t._version) for t in sources + ) + if source_state == self._fc12_source_state: + return + fc1, fc1_sf, fc2, fc2_sf = sources + if fc1.dtype != torch.float8_e4m3fn or fc2.dtype != torch.float8_e4m3fn: + raise ValueError("Shared FC12 requires loaded E4M3 weights") + if fc1.ndim != 2 or fc2.ndim != 2 or not fc1.is_contiguous() or not fc2.is_contiguous(): + raise ValueError("Shared FC12 weights must be contiguous matrices") + intermediate = fc2.shape[1] + if fc1.shape != (2 * intermediate, self.hidden_size) or fc2.shape[0] != self.hidden_size: + raise ValueError("Shared FC12 projection shapes do not describe the same TP shard") + if ( + intermediate <= 0 + or self.hidden_size <= 0 + or intermediate % 128 + or self.hidden_size % 128 + ): + raise ValueError( + "Shared FC12 hidden and local intermediate sizes must be divisible by 128" + ) + if any(t.device != fc1.device for t in sources): + raise ValueError("Shared FC12 weights and scales must be on the same device") + logical_fc1_sf = _unpack_r128c4(fc1_sf, 2 * intermediate, self.hidden_size // 32) + _unpack_r128c4(fc2_sf, self.hidden_size, intermediate // 32) + # Byte views keep packing independent of PyTorch FP8 indexing support. + packed_fc1 = _interleave_gate_up(fc1.view(torch.uint8)).view(fc1.dtype).unsqueeze(0) + packed_fc1_sf = _pack_r128c4(_interleave_gate_up(logical_fc1_sf)) + self._fc12_fc1_weight = packed_fc1 + self._fc12_fc1_weight_sf = packed_fc1_sf + self._fc12_fc2_weight = fc2.unsqueeze(0) + self._fc12_fc2_weight_sf = fc2_sf.unsqueeze(0) + self._fc12_source_state = source_state + + def forward( + self, + x: torch.Tensor, + all_rank_num_tokens: list[int] | None = None, + final_all_reduce_params: AllReduceParams | None = None, + lora_params: dict | None = None, + **kwargs, + ) -> torch.Tensor: + if lora_params or kwargs: + raise ValueError("Shared FC12 does not support LoRA or extra forward options") + if final_all_reduce_params is not None: + raise ValueError("Shared FC12 reduction belongs to the enclosing MoE module") + if not isinstance(x, torch.Tensor) or x.dtype != torch.bfloat16 or not x.is_cuda: + raise ValueError("Shared FC12 input must be a CUDA BF16 tensor") + if x.ndim != 2 or x.shape[1] != self.hidden_size or not x.is_contiguous(): + raise ValueError("Shared FC12 input must be contiguous [tokens, hidden_size]") + if x.shape[0] > self._fc12_max_tokens: + raise ValueError("Shared FC12 input exceeds configured max_num_tokens") + self.prepare_fc12_weights() + if x.device != self._fc12_fc1_weight.device: + raise ValueError("Shared FC12 input and weights must be on the same device") + from ..cute_dsl_kernels.megamoe_shared_fc12 import run_shared_fc12 + + return run_shared_fc12( + x, + self._fc12_fc1_weight, + self._fc12_fc1_weight_sf, + self._fc12_fc2_weight, + self._fc12_fc2_weight_sf, + swiglu_limit=self.swiglu_limit, + sm_count=self.sm_count, + max_tokens=self._fc12_max_tokens, + ) diff --git a/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py b/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py index 475619e3582b..fbe26c7c972f 100644 --- a/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py +++ b/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py @@ -35,6 +35,7 @@ import dataclasses import functools +import math import weakref from typing import Any, List, Optional, Tuple @@ -58,6 +59,9 @@ from tensorrt_llm.logger import logger from tensorrt_llm.math_utils import ceil_div, pad_up +# Default tactics validated once when the runtime becomes available. +_PREVALIDATED_TACTICS: List[Tuple] = [] + _NVFP4_BLOCK_SIZE = 16 _SF_PADDING_BLOCK = 128 _SUPPORTED_MMA_TILE_M = (128, 256) @@ -71,6 +75,7 @@ "enumerate_megamoe_candidate_tactics", "megamoe_activation_sf_bytes_per_row", "validate_megamoe_tactic", + "synthesize_profiling_topk", ] # Set to ``True`` if every symbol the op registration needs imports @@ -134,7 +139,7 @@ def _megamoe_fence_key_live(fkey: Tuple[int, int, int, str]) -> bool: # --------------------------------------------------------------------------- -# Tactic representation (v4: 10-tuple perf knobs) +# Tactic representation (v5: 11-tuple perf knobs) # --------------------------------------------------------------------------- # # A tactic is a tuple of JSON-friendly primitives (lists / ints / bools / @@ -151,7 +156,8 @@ def _megamoe_fence_key_live(fkey: Tuple[int, int, int, str]) -> bool: # use_bulk_fc2_store, # bool (epi -> peer UBLK; non-epi -> local TMA) # fc2_tma_stages, # optional staged-bulk depth # flag_batch, # int >= 1 (standalone_warps requires == 1) -# epi_flag_batch) # (int, int) fc1/fc2 done-counter publish batch +# epi_flag_batch, # (int, int) fc1/fc2 done-counter publish batch +# token_back_ready_granularity) # "expert" | "token_tile" # # Derived / out-of-tuple: # use_2cta_instrs = (mma_tiler_mnk[0] == 256), derived in _build_kernel. @@ -160,8 +166,11 @@ def _megamoe_fence_key_live(fkey: Tuple[int, int, int, str]) -> bool: # Tuple wrapping makes the tactic hashable, which AutoTuner needs for the # tactics cache. -_TACTIC_LEN = 10 +_TACTIC_LEN = 11 +_TACTIC_LEN_V4 = 10 _LEGACY_TACTIC_LEN = 8 +_DEFAULT_TOKEN_BACK_READY_GRANULARITY = "expert" # nosec B105 -- scheduler mode name, not a credential +_TOKEN_BACK_READY_GRANULARITIES = ("expert", "token_tile") # Kernel-side ceiling: ``flag_batch`` is hard-checked ``[1, 32]`` and # ``epi_flag_batch`` entries are SILENTLY clamped to ``[1, 32]``; reject > 32 @@ -169,14 +178,105 @@ def _megamoe_fence_key_live(fkey: Tuple[int, int, int, str]) -> bool: _FLAG_BATCH_MAX = 32 +_AUTOTUNE_PL_SEED = 0xA107 + + +def _autotune_pl_alpha(helper_expert_count: int = 0) -> float: + """Use the balanced-rank power-law workload only for load-balanced tuning.""" + return 0.8 if helper_expert_count > 0 else 0.0 + + +def synthesize_profiling_topk( + *, + num_tokens: int, + num_topk: int, + num_experts_per_rank: int, + world_size: int, + alpha: float, + device: torch.device, + source_rank: int = 0, +) -> torch.Tensor: + """Build physical routes with equal rank totals and ascending local load. + + Every source rank has T tokens and T*K routes. Rotating the destination + round-robin by source_rank*T*K makes the aggregate receive count exactly + T*K on every destination, including small buckets where EP does not + divide T*K. A single source's destination totals differ by at most one. + + Within each destination, sample with weights (E, ..., 1)**(-alpha). + Sampling is without replacement only within a token and destination; + equal local expert IDs on different ranks are distinct physical experts. + Relabel by each destination's realized counts so its expert load is + non-decreasing, including after summing all source ranks. Integer counts + fluctuate around the power-law weights; rank totals are exact. + + alpha=0 retains the OFF uniform round-robin byte for byte. This helper + runs only during profiling setup, outside candidate timing. + """ + tokens, topk = int(num_tokens), int(num_topk) + experts, ranks = int(num_experts_per_rank), int(world_size) + if tokens < 0 or experts <= 0 or ranks <= 0 or not 0 < topk <= experts * ranks: + raise ValueError("Invalid MegaMoE profiling route dimensions") + if not 0 <= source_rank < ranks or not math.isfinite(alpha) or alpha < 0: + raise ValueError("Invalid MegaMoE profiling source rank or exponent") + flat = torch.arange(tokens * topk, dtype=torch.long, device=device) + if alpha == 0: + return (flat % (experts * ranks)).view(tokens, topk) + destination = ((flat + source_rank * tokens * topk) % ranks).view(tokens, topk) + local_draws = (topk + ranks - 1) // ranks + log_weights = -alpha * torch.arange(experts, 0, -1, dtype=torch.float32, device=device).log() + generator = torch.Generator(device=device) + generator.manual_seed(_AUTOTUNE_PL_SEED) + noise = torch.rand((tokens, ranks, experts), generator=generator, device=device).clamp_( + min=torch.finfo(torch.float32).tiny + ) + gumbel = -torch.log(-torch.log(noise)) + candidates = torch.topk(log_weights + gumbel, local_draws, dim=-1).indices + local = candidates[ + torch.arange(tokens, device=device).unsqueeze(1), + destination, + torch.arange(topk, device=device) // ranks, + ] + counts = torch.bincount( + (destination * experts + local).reshape(-1), minlength=ranks * experts + ).view(ranks, experts) + order = torch.argsort(counts, dim=1, stable=True) + relabel = torch.empty_like(order) + relabel.scatter_(1, order, torch.arange(experts, device=device).expand(ranks, experts)) + return destination * experts + relabel[destination, local] + + +def _token_back_ready_granularity_choices(*, load_balance: bool = False) -> tuple[str, ...]: + """Compare both readiness modes only for load-balanced tuning.""" + return _TOKEN_BACK_READY_GRANULARITIES if load_balance else ("expert",) + + +def _megamoe_tuning_op_name(helper_expert_count: int) -> str: + """Separate timing caches without changing the registered op or kernel ABI.""" + base = "trtllm::cute_dsl_megamoe_nvfp4_blackwell" + load_balance = helper_expert_count > 0 + alpha = _autotune_pl_alpha(helper_expert_count) + granularities = _token_back_ready_granularity_choices(load_balance=load_balance) + if not load_balance and alpha == 0 and granularities == ("expert",): + return base + return ( + f"{base}:profile_v2:load_balance={int(load_balance)}:alpha={alpha!r}" + f":ready={','.join(granularities)}" + ) + + def _unpack_tactic(tactic: Tuple) -> Tuple: - """Return the canonical 10 fields, accepting the legacy 8-tuple. + """Return the canonical 11 fields, accepting the v4 10- and v3 8-tuples. Legacy tactics map ``static`` to grid-stride work IDs and preserve their grouped scheduling hint. Their single-stage bulk store is represented explicitly so old autotuner cache entries keep identical code generation. + Pre-v5 tactics take ``token_back_ready_granularity="expert"`` -- the only + value this op compiled before the field entered the tuple, so a cache + written by an older build keeps its exact meaning. Validation remains centralized in :func:`validate_megamoe_tactic`. """ + token_back_ready_granularity = _DEFAULT_TOKEN_BACK_READY_GRANULARITY if len(tactic) == _LEGACY_TACTIC_LEN: ( mma_tiler, @@ -200,6 +300,19 @@ def _unpack_tactic(tactic: Tuple) -> Tuple: f"'atomic_counter', got {load_balance_mode!r}." ) fc2_tma_stages = 1 if use_bulk_fc2_store else None + elif len(tactic) == _TACTIC_LEN_V4: + ( + mma_tiler, + cluster_shape, + fallback_cluster_shape, + schedule_policy, + work_id_mode, + token_back_mode, + use_bulk_fc2_store, + fc2_tma_stages, + flag_batch, + epi_flag_batch, + ) = tactic else: ( mma_tiler, @@ -212,6 +325,7 @@ def _unpack_tactic(tactic: Tuple) -> Tuple: fc2_tma_stages, flag_batch, epi_flag_batch, + token_back_ready_granularity, ) = tactic return ( @@ -225,14 +339,60 @@ def _unpack_tactic(tactic: Tuple) -> Tuple: fc2_tma_stages, flag_batch, epi_flag_batch, + token_back_ready_granularity, ) +# Reuse default tactic objects so validation can be memoized by identity. +# These shared constants must not be mutated. Value-only memoization could +# conflate bool with int and bypass tactic type checks. +_DEFAULT_TACTIC_DECODE: Tuple = ( + [256, 128, 256], + [2, 1, 1], + 512, + "static", + "epi_warps", + True, + 1, + (1, 1), +) +_DEFAULT_TACTIC_MID: Tuple = ( + [256, 256, 256], + [2, 1, 1], + 512, + "static", + "epi_warps", + True, + 4, + (1, 1), +) +_DEFAULT_TACTIC_PREFILL_STATIC: Tuple = ( + [256, 256, 256], + [2, 1, 1], + 512, + "static", + "reuse_dispatch_warps", + False, + 8, + (2, 4), +) +_DEFAULT_TACTIC_PREFILL_ATOMIC: Tuple = ( + [256, 256, 256], + [2, 1, 1], + 512, + "atomic_counter", + "reuse_dispatch_warps", + False, + 8, + (2, 4), +) + + def default_megamoe_tactic(num_tokens: int) -> Tuple: """Deterministic token-bucket fallback when no shape-tuned default exists; never profiled by the autotuner.""" if num_tokens <= 1024: - # decode winner: N128 only helps epi_warps + bulk. + # Deterministic fallback for the smallest token bucket. return ( [256, 128, 256], [2, 1, 1], @@ -244,6 +404,7 @@ def default_megamoe_tactic(num_tokens: int) -> Tuple: 1, 1, (1, 1), + _DEFAULT_TOKEN_BACK_READY_GRANULARITY, ) if num_tokens <= 8192: return ( @@ -257,8 +418,9 @@ def default_megamoe_tactic(num_tokens: int) -> Tuple: 1, 4, (1, 1), + _DEFAULT_TOKEN_BACK_READY_GRANULARITY, ) - # prefill: atomic_counter only helps the very large tail (>=16384). + # Use the atomic work distributor for the upper token bucket. return ( [256, 256, 256], [2, 1, 1], @@ -270,13 +432,13 @@ def default_megamoe_tactic(num_tokens: int) -> Tuple: None, 8, (2, 4), + _DEFAULT_TOKEN_BACK_READY_GRANULARITY, ) -# Static defaults for the SM107 DSv4 Pro DEP4 and DEP8 shapes. -# Routing distribution is not a -# rank-identical runtime input, so each bucket uses one static tactic. The full -# runner key prevents applying the table beyond the measured shape and codegen modes. +# Static defaults for the explicitly keyed SM107 configurations. Routing is not +# rank-identical, so each bucket uses one deterministic tactic. The complete +# runner key prevents applying an entry to a different shape or codegen mode. _SM107_DSV4_PRO_DEFAULT_TACTICS: dict[Tuple, Tuple] = { ( 107, @@ -296,6 +458,7 @@ def default_megamoe_tactic(num_tokens: int) -> Tuple: "bf16", None, None, + 0, ): ( list(mma_tiler), [4, 1, 1], @@ -524,6 +687,7 @@ def _megamoe_problem_key( combine_format: str, swiglu_alpha: Optional[float] = None, swiglu_beta: Optional[float] = None, + helper_expert_count: int = 0, ) -> Tuple: return ( int(sm_version), @@ -543,6 +707,8 @@ def _megamoe_problem_key( str(combine_format), None if swiglu_alpha is None else float(swiglu_alpha), None if swiglu_beta is None else float(swiglu_beta), + # The helper count changes generated code and must be part of its cache key. + int(helper_expert_count), ) @@ -557,19 +723,19 @@ def _default_megamoe_tactic_for_problem( shape_tuned_tactic: Optional[Tuple] = None, ) -> Tuple: if combine_format != "bf16": - return _MEGAMOE_NONBULK_STANDALONE_TACTIC + return _unpack_tactic(_MEGAMOE_NONBULK_STANDALONE_TACTIC) bucket = _sm107_tactic_bucket(max_tokens_per_rank) if sm_version == 107 else None if bucket is not None: expected_in_kernel_fc2_reduce = _sm107_use_in_kernel_fc2_reduce( max_tokens_per_rank, apply_topk_in_fc1 ) if in_kernel_fc2_reduce != expected_in_kernel_fc2_reduce: - return _MEGAMOE_NONBULK_STANDALONE_TACTIC - return _SM107_GENPHASE_TACTICS[bucket] + return _unpack_tactic(_MEGAMOE_NONBULK_STANDALONE_TACTIC) + return _unpack_tactic(_SM107_GENPHASE_TACTICS[bucket]) if in_kernel_fc2_reduce: - return _MEGAMOE_NONBULK_STANDALONE_TACTIC + return _unpack_tactic(_MEGAMOE_NONBULK_STANDALONE_TACTIC) if shape_tuned_tactic is not None: - return shape_tuned_tactic + return _unpack_tactic(shape_tuned_tactic) return default_megamoe_tactic(num_tokens) @@ -587,14 +753,16 @@ def _default_megamoe_tactic_for_key(problem_key: Tuple, *, num_tokens: int) -> T def _megamoe_default_tactic_candidates(tactic: Tuple) -> Tuple[Tuple, ...]: """Return the heuristic tactic followed by shape-safe fallbacks.""" + tactic = _unpack_tactic(tactic) candidates = [tactic] - schedule_policy = _unpack_tactic(tactic)[3] + schedule_policy = tactic[3] if schedule_policy[0] == "phase_interleave" and schedule_policy[1] is not None: auto_hint_tactic = list(tactic) auto_hint_tactic[3] = ("phase_interleave", None) candidates.append(tuple(auto_hint_tactic)) - if _MEGAMOE_NONBULK_STANDALONE_TACTIC not in candidates: - candidates.append(_MEGAMOE_NONBULK_STANDALONE_TACTIC) + nonbulk_tactic = _unpack_tactic(_MEGAMOE_NONBULK_STANDALONE_TACTIC) + if nonbulk_tactic not in candidates: + candidates.append(nonbulk_tactic) return tuple(candidates) @@ -713,13 +881,16 @@ def _construct_megamoe_kernel( """Construct the active architecture kernel and report a narrow gen rejection.""" from cutlass.cute.nvgpu import OperandMajorMode - from tensorrt_llm._torch.moe.kernels.cutedsl_megamoe import ( + from tensorrt_llm._torch.moe.kernels.cutedsl_megamoe.api import ImplDesc, ProblemDesc + from tensorrt_llm._torch.moe.kernels.cutedsl_megamoe.quant_def import ( + CombineFormat, + QuantKind, + ) + from tensorrt_llm._torch.moe.kernels.megamoe_kernels import ( BlackwellInferenceMegaMoE, RubinInferenceGenphaseMegaMoE, RubinInferenceMegaMoE, ) - from tensorrt_llm._torch.moe.kernels.cutedsl_megamoe.api import ImplDesc, ProblemDesc - from tensorrt_llm._torch.moe.kernels.cutedsl_megamoe.quant_def import CombineFormat, QuantKind local_experts, intermediate_gateup_size, hidden_size = common["static_expert_shape"] problem_desc = ProblemDesc( @@ -754,12 +925,15 @@ def _construct_megamoe_kernel( "work_id_mode": str(common["work_id_mode"]), "fc2_use_bulk": not bool(common["non_ubulk_fc2_store"]), "epi_flag_batches": tuple(common["epi_flag_batch"]), + # Kernel construction and READY arguments must use the same helper count. + "helper_expert_count": int(common.get("helper_expert_count", 0)), "launch_cluster_count": int(common["launch_cluster_count"]), "fallback_cluster_shape_mn": common["fallback_cluster_shape_mn"], "preferred_cluster_count": common["preferred_cluster_count"], "fallback_cluster_count": common["fallback_cluster_count"], "token_in_flag_batch": int(common["flag_batch"]), "token_back_mode": str(common["token_back_mode"]), + "token_back_ready_granularity": str(common.get("token_back_ready_granularity", "expert")), "reduce_topk_in_kernel": bool(common["in_kernel_fc2_reduce"]), } if sm_version == 107: @@ -775,15 +949,25 @@ def _construct_megamoe_kernel( impl_values["fc2_tma_stages"] = int(common["fc2_tma_stages"]) impl_desc = ImplDesc(impl_values) + def construct(kernel_cls): + # ProblemDesc.validate ignores undeclared keys, so a kernel without the SwiGLU + # constants would silently run plain SwiGLU. + if ( + common["swiglu_alpha"] is not None + and "swiglu_alpha" not in kernel_cls.problem_desc_require() + ): + raise NotImplementedError(f"{kernel_cls.__name__} does not support swiglu_alpha.") + return kernel_cls(problem_desc, impl_desc) + bucket = _sm107_tactic_bucket(int(common["max_tokens_per_rank"])) if sm_version == 107 else None rejection = None if bucket is not None and bucket not in _SM107_GENERIC_KERNEL_BUCKETS: try: - return RubinInferenceGenphaseMegaMoE(problem_desc, impl_desc), None + return construct(RubinInferenceGenphaseMegaMoE), None except (ValueError, TypeError, NotImplementedError) as exc: rejection = exc kernel_cls = RubinInferenceMegaMoE if sm_version == 107 else BlackwellInferenceMegaMoE - return kernel_cls(problem_desc, impl_desc), rejection + return construct(kernel_cls), rejection def _megamoe_barrier_reset_slices(kernel) -> Tuple[int, int, int, int, str]: @@ -822,11 +1006,12 @@ def validate_megamoe_tactic(tactic: Tuple, sm_version: int = 100) -> None: """ if (not isinstance(tactic, tuple)) or len(tactic) not in { _TACTIC_LEN, + _TACTIC_LEN_V4, _LEGACY_TACTIC_LEN, }: raise ValueError( f"MegaMoE tactic must be a {_TACTIC_LEN}-tuple (or legacy " - f"{_LEGACY_TACTIC_LEN}-tuple), got " + f"{_TACTIC_LEN_V4}- / {_LEGACY_TACTIC_LEN}-tuple), got " f"{type(tactic).__name__} len={len(tactic) if isinstance(tactic, tuple) else 'NA'}={tactic!r}" ) is_legacy_tactic = len(tactic) == _LEGACY_TACTIC_LEN @@ -841,6 +1026,7 @@ def validate_megamoe_tactic(tactic: Tuple, sm_version: int = 100) -> None: fc2_tma_stages, flag_batch, epi_flag_batch, + token_back_ready_granularity, ) = _unpack_tactic(tactic) if (not isinstance(mma_tiler, (list, tuple))) or len(mma_tiler) != 3: @@ -1001,10 +1187,89 @@ def validate_cluster_shape(shape, field_name: str) -> None: f"(epilogue clamp range), got {epi_flag_batch!r}." ) + if token_back_ready_granularity not in _TOKEN_BACK_READY_GRANULARITIES: + raise ValueError( + f"token_back_ready_granularity must be one of " + f"{_TOKEN_BACK_READY_GRANULARITIES}, got {token_back_ready_granularity!r}." + ) + if token_back_ready_granularity == "token_tile": # nosec B105 -- scheduler mode name, not a credential + # Mirrors the kernel's own guard (kernel_src/*/inference/mega/ + # block_scaled_swap_ab_mega_moe_kernel.py, _validate_geometry). Checking + # it HERE is the point of putting the granularity in the tuple: the + # kernel raises at build time, which is a crash, while enumeration wraps + # every candidate in this function and simply drops what fails. The + # kernel-side copy stays as the backstop for hand-supplied tactics. + # + # Two of the kernel's nine conditions are not tactic-local and are not + # re-checked here: + # quant_kind == nvfp4 -- this op's ProblemDesc hardcodes nvfp4. + # not reduce_topk_in_kernel -- implied transitively: token_tile + # requires use_bulk_fc2_store, and get_valid_tactics already drops + # every bulk candidate when in_kernel_fc2_reduce is set. + unmet = [] + if work_id_mode != "atomic_counter": + unmet.append(f"work_id_mode={work_id_mode!r} (need 'atomic_counter')") + if token_back_mode not in ("standalone_warps", "reuse_dispatch_warps"): + unmet.append( + f"token_back_mode={token_back_mode!r} " + f"(need 'standalone_warps' or 'reuse_dispatch_warps')" + ) + if mma_tiler[0] != 256: + unmet.append(f"mma_tiler_mnk[0]={mma_tiler[0]} (need 256, i.e. use_2cta_instrs)") + if not use_bulk_fc2_store: + unmet.append("use_bulk_fc2_store=False (need True, FC2 TMA stores)") + if unmet: + raise ValueError( + "token_back_ready_granularity='token_tile' is unmet by: " + "; ".join(unmet) + ) + + +def _expand_megamoe_ready_tactics( + candidates: list[Tuple], sm_version: int, *, load_balance: bool +) -> list[Tuple]: + """Return canonical tactics for every legal token-back READY mode.""" + choices = _token_back_ready_granularity_choices(load_balance=load_balance) + if sm_version == 107 and choices != ("expert",): + # At every token bucket include a legal local-TMA/atomic-counter + # base; the historical candidates otherwise exclude token_tile. + for cluster, fallback, policy, mode, flag, epi in ( + ([2, 1, 1], None, ("grouped", 512), "reuse_dispatch_warps", 4, (1, 1)), + ([2, 1, 1], None, ("grouped", 512), "reuse_dispatch_warps", 4, (2, 4)), + ([2, 1, 1], None, ("grouped", 512), "standalone_warps", 1, (2, 4)), + ([4, 1, 1], [2, 1, 1], ("phase_interleave", None), "reuse_dispatch_warps", 4, (2, 4)), + ): + candidates.append( + ( + [256, 256, 256], + cluster, + fallback, + policy, + "atomic_counter", + mode, + True, + 2, + flag, + epi, + ) + ) + expanded = [] + for base in candidates: + for granularity in choices: + tactic = (*_unpack_tactic(base)[:_TACTIC_LEN_V4], granularity) + try: + validate_megamoe_tactic(tactic, sm_version=sm_version) + except ValueError: + continue + if tactic not in expanded: + expanded.append(tactic) + return expanded + def enumerate_megamoe_candidate_tactics( num_tokens: int, sm_version: int = 100, + *, + load_balance: bool = False, ) -> List[Tuple]: """Return the curated candidate tactic list for the current token bucket. @@ -1093,7 +1358,7 @@ def enumerate_megamoe_candidate_tactics( ) validate_megamoe_tactic(mixed, sm_version=sm_version) candidates.append(mixed) - return candidates + return _expand_megamoe_ready_tactics(candidates, sm_version, load_balance=load_balance) if sm_version not in (100, 103): raise ValueError( f"MegaMoE tactic enumeration supports SM100, SM103, or SM107; got SM{sm_version}." @@ -1125,11 +1390,8 @@ def enumerate_megamoe_candidate_tactics( continue candidates.append(tactic) - # These additions are intentionally curated instead of crossed with every - # legacy axis: staged FC2 UBLK and mixed-CGA/phase scheduling are expensive - # to compile and autotuning is opt-in. These entries expose the upstream - # performance mechanisms while preserving the existing deterministic - # production defaults. + # Add a bounded set of valid staged and mixed-scheduling candidates rather + # than crossing every independent tactic axis. if num_tokens >= 2048: for fc2_tma_stages in (2, 4): tactic = ( @@ -1146,8 +1408,7 @@ def enumerate_megamoe_candidate_tactics( ) validate_megamoe_tactic(tactic, sm_version=sm_version) candidates.append(tactic) - # Local TMA stage 2 was the only consistently useful depth in B200 - # A/B tests: stage 1 tied it while stage 4 regressed from extra SMEM. + # Include the resource-balanced local TMA configuration. tactic = ( [256, 256, 256], [2, 1, 1], @@ -1177,7 +1438,7 @@ def enumerate_megamoe_candidate_tactics( ) validate_megamoe_tactic(tactic, sm_version=sm_version) candidates.append(tactic) - return candidates + return _expand_megamoe_ready_tactics(candidates, sm_version, load_balance=load_balance) if IS_CUTLASS_DSL_AVAILABLE: @@ -1200,15 +1461,15 @@ def enumerate_megamoe_candidate_tactics( except ImportError: from cuda import cuda - from tensorrt_llm._torch.moe.kernels.cutedsl_megamoe import ( # noqa: F401 + from tensorrt_llm._torch.moe.kernels.cutedsl_megamoe.communication.nvlink_domain.symmetric_buffer import ( + SymmetricBufferHost, # noqa: F401 + ) + from tensorrt_llm._torch.moe.kernels.megamoe_kernels import ( # noqa: F401 BlackwellInferenceMegaMoE, RubinInferenceGenphaseMegaMoE, RubinInferenceLocalMegaMoE, RubinInferenceMegaMoE, ) - from tensorrt_llm._torch.moe.kernels.cutedsl_megamoe.communication.nvlink_domain.symmetric_buffer import ( - SymmetricBufferHost, # noqa: F401 - ) IS_MEGAMOE_OP_AVAILABLE = True except Exception as _megamoe_import_err: # pragma: no cover - env-specific @@ -1227,6 +1488,27 @@ def enumerate_megamoe_candidate_tactics( if IS_MEGAMOE_OP_AVAILABLE: + # Validate default tactics before enabling identity-based validation bypass. + # If validation fails, retain normal per-call checks rather than failing import. + for _prevalidate_t in ( + _DEFAULT_TACTIC_DECODE, + _DEFAULT_TACTIC_MID, + _DEFAULT_TACTIC_PREFILL_STATIC, + _DEFAULT_TACTIC_PREFILL_ATOMIC, + _MEGAMOE_NONBULK_STANDALONE_TACTIC, + ): + try: + validate_megamoe_tactic(_prevalidate_t) + except Exception as _prevalidate_err: # pragma: no cover - partial installation + logger.debug( + "[MegaMoE] tactic prevalidation skipped (%s); falling back to " + "per-launch validation.", + _prevalidate_err, + ) + _PREVALIDATED_TACTICS.clear() + break + _PREVALIDATED_TACTICS.append(_prevalidate_t) + # ----- Local workspace cache -------------------------------------------- # # ``local_workspace`` is per-rank CUDA-only and sized by @@ -1243,6 +1525,10 @@ def enumerate_megamoe_candidate_tactics( # freed when the next one allocates (only the winner is reused). _MEGAMOE_TUNING_WORKSPACE_KEYS: set = set() + # Cache runners by constructor inputs only outside tactic autotuning. + # Reset profiling scratch on reuse and overwrite READY inputs for every call; + _MEGAMOE_RUNNER_CACHE: dict = {} + @functools.lru_cache(maxsize=1) def _cute_launch_helpers(): """Import the cutlass-dsl launch surface once per process and build the @@ -1752,6 +2038,7 @@ def build(candidate: Tuple): fc2_tma_stages, flag_batch, epi_flag_batch, + token_back_ready_granularity, ) = _unpack_tactic(candidate) mma_tiler = tuple(mma_tiler) ( @@ -1796,6 +2083,7 @@ def build(candidate: Tuple): fc2_output_dtype=cutlass.BFloat16, in_kernel_fc2_reduce=bool(in_kernel_fc2_reduce), token_back_mode=str(token_back_mode), + token_back_ready_granularity=str(token_back_ready_granularity), non_ubulk_fc2_store=(not bool(use_bulk_fc2_store)), flag_batch=int(flag_batch), epi_flag_batch=tuple(epi_flag_batch), @@ -1877,6 +2165,7 @@ def __init__( in_kernel_fc2_reduce: bool = False, combine_format: str = "bf16", tactic_autotune: bool = False, + helper_expert_count: int = 0, ) -> None: super().__init__() if (sm_version := get_sm_version()) not in (100, 103, 107): @@ -1929,6 +2218,15 @@ def __init__( # generated kernel, so deliberately EXCLUDED from unique_id / # _tactic_cache_key: opted-in and opted-out runs share caches. self.tactic_autotune = bool(tactic_autotune) + # ``S``. CODEGEN-CHANGING: a helper-bearing build carries the READY + # gate and a static ``H`` boundary in the sched warp, so it MUST NOT + # share a compile-cache entry with the ungated baseline. Hence it + # is in unique_id() (which _tactic_cache_key wraps). + self.helper_expert_count = int(helper_expert_count) + # Per-launch READY runtime bundle (device pointer + generation). + # Set by the op right after construction; never a codegen input. + self._helper_ready_flags = None + self._helper_ready_generation = 0 def unique_id(self): # local_rank is intentionally excluded: every EP rank must run the @@ -1955,6 +2253,7 @@ def unique_id(self): situ_linear_beta=self.situ_linear_beta, in_kernel_fc2_reduce=self.in_kernel_fc2_reduce, combine_format=self.combine_format, + helper_expert_count=self.helper_expert_count, ) def get_valid_tactics( @@ -1965,7 +2264,11 @@ def get_valid_tactics( ) -> List[Tuple]: del profile, kwargs num_tokens = int(inputs[0].shape[0]) - candidates = enumerate_megamoe_candidate_tactics(num_tokens, sm_version=self.sm_version) + candidates = enumerate_megamoe_candidate_tactics( + num_tokens, + sm_version=self.sm_version, + load_balance=self.helper_expert_count > 0, + ) # Form-B remains non-bulk. FP4 only rejects peer UBLK; local TMA # supports its packed payload. Fields 5/6 are token-back mode and # use_bulk_fc2_store respectively. @@ -1988,17 +2291,22 @@ def get_valid_tactics( if default_tactic not in candidates: candidates.append(default_tactic) default_tactic = _SM107_DSV4_PRO_DEFAULT_TACTICS.get(self.unique_id()) - if default_tactic is not None and default_tactic not in candidates: - candidates.append(default_tactic) - return candidates + if default_tactic is not None: + default_tactic = _unpack_tactic(default_tactic) + if default_tactic not in candidates: + candidates.append(default_tactic) + choices = _token_back_ready_granularity_choices( + load_balance=self.helper_expert_count > 0 + ) + return [t for t in candidates if _unpack_tactic(t)[10] in choices] def _autotuner_inputs_pre_hook(self, inputs: List[torch.Tensor]) -> List[torch.Tensor]: """Sanitize ONLY the autotuner-regenerated fake inputs. AutoTuner rebuilds inputs 0-3 and 11; the static inputs (4-10) are the caller's REAL weights/scales, passed by reference -- filling - those would clobber them. topk_idx becomes a valid round-robin - (random ints index OOB); SF / weights are filled with 1.0, NOT 0 + those would clobber them. ON topk_idx uses rank-balanced power-law + physical routes; OFF keeps round-robin. SF / weights use 1.0, NOT 0 (SF==0 degenerates the GEMMs and skews tactic timing). """ # Runs ONLY on a real profiling MISS (a tuning-mode cache HIT @@ -2026,15 +2334,19 @@ def _autotuner_inputs_pre_hook(self, inputs: List[torch.Tensor]) -> List[torch.T topk_idx = inputs[2] if isinstance(topk_idx, torch.Tensor) and topk_idx.dim() == 2: T, K = topk_idx.shape - valid = ( - torch.arange( - T * K, - dtype=topk_idx.dtype, - device=topk_idx.device, - ) - % total_experts - ).view(T, K) - topk_idx.copy_(valid) + alpha = _autotune_pl_alpha(self.helper_expert_count) + valid = synthesize_profiling_topk( + num_tokens=T, + num_topk=K, + num_experts_per_rank=self.num_experts_per_rank, + world_size=self.world_size, + source_rank=self.local_rank, + alpha=alpha, + device=topk_idx.device, + ) + topk_idx.copy_(valid.to(topk_idx.dtype)) + if self.local_rank == 0 and alpha > 0: + logger.debug("[MegaMoE] helper-bearing autotune workload prepared") scratch = self._profiling_scratch if scratch is not None: @@ -2143,6 +2455,7 @@ def _build_kernel_once(self, tactic: Tuple): fc2_tma_stages, flag_batch, epi_flag_batch, + token_back_ready_granularity, ) = _unpack_tactic(tactic) ( launch_cluster_count, @@ -2186,6 +2499,7 @@ def _build_kernel_once(self, tactic: Tuple): fc2_output_dtype=cutlass.BFloat16, in_kernel_fc2_reduce=self.in_kernel_fc2_reduce, token_back_mode=str(token_back_mode), + token_back_ready_granularity=str(token_back_ready_granularity), non_ubulk_fc2_store=(not bool(use_bulk_fc2_store)), flag_batch=int(flag_batch), epi_flag_batch=tuple(epi_flag_batch), @@ -2195,6 +2509,7 @@ def _build_kernel_once(self, tactic: Tuple): gate_up_clamp=self.gate_up_clamp, situ_beta=self.situ_beta, situ_linear_beta=self.situ_linear_beta, + helper_expert_count=self.helper_expert_count, **_LOCKED_KERNEL_KWARGS, ) kernel, rejection = _construct_megamoe_kernel( @@ -2228,6 +2543,7 @@ def _tactic_cache_key(self, tactic: Tuple) -> Tuple: fc2_tma_stages, flag_batch, epi_flag_batch, + token_back_ready_granularity, ) = _unpack_tactic(tactic) launch_config = _launch_cluster_configuration( cluster_shape, @@ -2246,6 +2562,7 @@ def _tactic_cache_key(self, tactic: Tuple) -> Tuple: fc2_tma_stages, int(flag_batch), tuple(epi_flag_batch), + str(token_back_ready_granularity), launch_config, ) @@ -2327,10 +2644,14 @@ def forward( # -- and a full zero wastes decode time). form-A overwrites. # Cached per FULL tactic: local workspace SIZE is tactic-dependent - # (see _tactic_cache_key). + # (see _tactic_cache_key). The key is computed HERE, not at the + # launch site further down: the local-workspace cache needs it + # first, and `tactic_t` is final from `_resolve_megamoe_kernel_tactic` + # above. + cache_key = self._tactic_cache_key(tactic_t) local_workspace = _get_or_alloc_local_workspace( kernel, - cache_key=self._tactic_cache_key(tactic_t), + cache_key=cache_key, device=activation.device, latch_in_tuning_mode=self.tactic_autotune, ) @@ -2496,6 +2817,37 @@ def forward( ), ) cache_key = self._tactic_cache_key(tactic_t) + # READY ABI tail: attached only when S > 0, i.e. rebalance is on. + # Contract: with S > 0 both values must be supplied. Fail loudly + # here rather than hand the kernel an unarmed gate. + if self.helper_expert_count > 0: + import cutlass + + if self._helper_ready_flags is None: + raise RuntimeError( + "MegaMoE-CuteDSL: helper_expert_count > 0 requires " + "hot_expert_weight_ready_flags (uint64[EP] terminal " + "flags); got None." + ) + if not (1 <= int(self._helper_ready_generation) < (1 << 63)): + raise RuntimeError( + "MegaMoE-CuteDSL: hot_expert_weight_ready_generation " + f"must be in [1, 2**63); got " + f"{self._helper_ready_generation}." + ) + if self._helper_ready_flags.numel() < self.world_size: + raise RuntimeError( + "MegaMoE-CuteDSL: hot_expert_weight_ready_flags must " + f"hold at least world_size={self.world_size} uint64 " + f"cells; got numel={self._helper_ready_flags.numel()}." + ) + runtime_kwargs["hot_expert_weight_ready_flags"] = _to_cute_ptr( + self._helper_ready_flags + ) + runtime_kwargs["hot_expert_weight_ready_generation"] = cutlass.Uint64( + int(self._helper_ready_generation) + ) + # unique_id() includes helper_expert_count, separating cached ABIs. cached = self.__class__.kernel_cache.get(cache_key) if cached is None: compiled = cute.compile(kernel, **runtime_kwargs) @@ -2505,6 +2857,89 @@ def forward( compiled(**runtime_kwargs) return combine_output + def _megamoe_get_runner( + world_size, + local_rank, + num_topk, + num_experts_per_rank, + hidden_size, + intermediate_size_per_partition, + expand_intermediate_size_per_partition, + max_tokens_per_rank, + output_dtype, + apply_topk_in_fc1, + swiglu_alpha, + swiglu_beta, + gate_up_clamp, + situ_beta, + situ_linear_beta, + in_kernel_fc2_reduce, + combine_format, + tactic_autotune, + helper_expert_count, + ): + """Reuse a runner by its complete constructor key. + + Reset profiling scratch on reuse; each op call supplies fresh READY inputs. + """ + args = ( + world_size, + local_rank, + num_topk, + num_experts_per_rank, + hidden_size, + intermediate_size_per_partition, + expand_intermediate_size_per_partition, + max_tokens_per_rank, + output_dtype, + apply_topk_in_fc1, + swiglu_alpha, + swiglu_beta, + gate_up_clamp, + situ_beta, + situ_linear_beta, + in_kernel_fc2_reduce, + combine_format, + tactic_autotune, + helper_expert_count, + ) + + def _make(): + return Sm100MegaMoENvfp4Runner( + world_size=world_size, + local_rank=local_rank, + num_topk=num_topk, + num_experts_per_rank=num_experts_per_rank, + hidden_size=hidden_size, + intermediate_size_per_partition=intermediate_size_per_partition, + expand_intermediate_size_per_partition=expand_intermediate_size_per_partition, + max_tokens_per_rank=max_tokens_per_rank, + output_dtype=output_dtype, + apply_topk_in_fc1=apply_topk_in_fc1, + swiglu_alpha=swiglu_alpha, + swiglu_beta=swiglu_beta, + gate_up_clamp=gate_up_clamp, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, + in_kernel_fc2_reduce=in_kernel_fc2_reduce, + combine_format=combine_format, + tactic_autotune=tactic_autotune, + helper_expert_count=helper_expert_count, + ) + + # Autotuning runners own profiling scratch and are not reused. + if tactic_autotune: + return _make() + runner = _MEGAMOE_RUNNER_CACHE.get(args) + if runner is None: + runner = _make() + _MEGAMOE_RUNNER_CACHE[args] = runner + else: + # Reset profiling scratch on reuse; the caller supplies fresh READY fields. + runner._profiling_scratch = None + runner._profiling_scratch_factory = None + return runner + # ----- torch op --------------------------------------------------------- @torch.library.custom_op( @@ -2545,6 +2980,14 @@ def cute_dsl_megamoe_nvfp4_blackwell( combine_format: str = "bf16", tactic_autotune: bool = False, num_tokens: int = -1, + # READY ABI tail. ``helper_expert_count`` (S) is codegen-time; + # ``hot_expert_weight_ready_flags`` is the local uint64[EP] terminal + # view published by the weight transport, ``..._generation`` the + # expected generation g in [1, 2**63). All three default to the + # unchanged S == 0 baseline. + helper_expert_count: int = 0, + hot_expert_weight_ready_flags: Optional[torch.Tensor] = None, + hot_expert_weight_ready_generation: int = 0, ) -> None: """Run the fused MegaMoE CuteDSL NVFP4 kernel. @@ -2586,26 +3029,31 @@ def cute_dsl_megamoe_nvfp4_blackwell( if num_tokens > 0: combine_output = combine_output[: min(num_tokens, combine_output.shape[0])] - runner = Sm100MegaMoENvfp4Runner( - world_size=world_size, - local_rank=local_rank, - num_topk=num_topk, - num_experts_per_rank=num_experts_per_rank, - hidden_size=hidden_size, - intermediate_size_per_partition=intermediate_size_per_partition, - expand_intermediate_size_per_partition=expand_intermediate_size_per_partition, - max_tokens_per_rank=max_tokens_per_rank, - output_dtype=combine_output.dtype, - apply_topk_in_fc1=apply_topk_in_fc1, - swiglu_alpha=swiglu_alpha, - swiglu_beta=swiglu_beta, - gate_up_clamp=gate_up_clamp, - situ_beta=act_alpha, - situ_linear_beta=act_beta, - in_kernel_fc2_reduce=in_kernel_fc2_reduce, - combine_format=combine_format, - tactic_autotune=tactic_autotune, + runner = _megamoe_get_runner( + world_size, + local_rank, + num_topk, + num_experts_per_rank, + hidden_size, + intermediate_size_per_partition, + expand_intermediate_size_per_partition, + max_tokens_per_rank, + combine_output.dtype, + apply_topk_in_fc1, + swiglu_alpha, + swiglu_beta, + gate_up_clamp, + act_alpha, + act_beta, + in_kernel_fc2_reduce, + combine_format, + tactic_autotune, + helper_expert_count, ) + # READY inputs belong to this launch. Never regenerate them or retain stale + # values across calls, even when the runner itself is cached. + runner._helper_ready_flags = hot_expert_weight_ready_flags + runner._helper_ready_generation = int(hot_expert_weight_ready_generation) inputs = [ activation, activation_sf, @@ -2626,12 +3074,17 @@ def cute_dsl_megamoe_nvfp4_blackwell( # materializes the multi-GiB scratch and MERGE-sweeps every # candidates. tactic=-1 guarantees deterministic default selection # with no tuning collectives, even inside a global autotune(). - runner( - inputs, - tactic=-1, - peer_offsets=peer_offsets, - shared_workspace=shared_workspace, - ) + try: + runner( + inputs, + tactic=-1, + peer_offsets=peer_offsets, + shared_workspace=shared_workspace, + ) + finally: + # Release the cached runner reference to this generation's terminal tensor. + # The next call supplies fresh READY inputs. + runner._helper_ready_flags = None return tuner = AutoTuner.get() # Opt-IN: in tuning mode the MERGE lockstep sweep runs (made safe by @@ -2670,7 +3123,7 @@ def cute_dsl_megamoe_nvfp4_blackwell( prof_shared_workspace = shared_workspace try: _, best_tactic = tuner.choose_one( - "trtllm::cute_dsl_megamoe_nvfp4_blackwell", + _megamoe_tuning_op_name(runner.helper_expert_count), [runner], runner.get_tuning_config(), inputs, @@ -2722,5 +3175,13 @@ def _( combine_format: str = "bf16", tactic_autotune: bool = False, num_tokens: int = -1, + # READY ABI tail. ``helper_expert_count`` (S) is codegen-time; + # ``hot_expert_weight_ready_flags`` is the local uint64[EP] terminal + # view published by the weight transport, ``..._generation`` the + # expected generation g in [1, 2**63). All three default to the + # unchanged S == 0 baseline. + helper_expert_count: int = 0, + hot_expert_weight_ready_flags: Optional[torch.Tensor] = None, + hot_expert_weight_ready_generation: int = 0, ) -> None: return None diff --git a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md index a02033e6794f..f1966bc21eb6 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -438,8 +438,10 @@ literal also selects, descend from the CUTLASS family and stay in | `KERNEL_INTERNALS.html` | Reference for the underlying DeepGEMM kernel layout | The exported CuteDSL kernel sources for `TrtllmCutedslMegaMoeNvfp4Impl` live under -`tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/`, preserving the upstream -package structure and `KernelClass` entry points. The runner constructs +`tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/`, vendored from upstream as +recorded in `3rdparty/vendor_sources.lock.yaml`; keep local edits out of that +directory. `moe/kernels/megamoe_kernels.py` exposes the `KernelClass` entry +points. The runner constructs `ProblemDesc` and `ImplDesc` and dispatches internally; the public backend remains `MEGAMOE_CUTEDSL`. Rubin's smaller token buckets use the genphase kernel when supported, otherwise the generic kernel. During warmup, the @@ -861,6 +863,34 @@ set wrong has a test that catches it: - `FusedCommMoEScheduler` calls `backend.quantize_input(...)` for every chunk including zero-token chunks (so peer ranks can cross the in-kernel NVLink barrier). Each fused-comm backend therefore MUST make `quantize_input` tolerate `x.shape[0] == 0` and return its own empty tensor layout; the scheduler does NOT synthesize backend-specific empty tensors. - Dynamic EPLB requires backend and quantization-method support; backends gate wrapper-level constraints in `validate_configurable_moe`. `MegaMoEDeepGemm` routes to slot IDs and migrates the transformed DG weights its quantization method registers (`num_slots % ep_size == 0`). `TrtllmCutedslMegaMoeNvfp4Impl` declares `supports_eplb=True`; its quantization method registers the MegaMoE-format derived params and `fc1_norm_const` with the load balancer alongside the raw NVFP4 family, so per-slot migration stays byte-consistent. - `FUSED_COMM` backends use `ignore_allreduce=False` for EPLB statistic update because the fused kernel AllReduces routing stats internally. +- Per-layer dynamic EPLB uses `MoeConfig.load_balancer.mode=per_layer_dynamic` with + `num_slots` greater than the model expert count. Backend resolution selects + `MEGAMOE_CUTEDSL` when no concrete backend is requested and rejects an + incompatible explicit backend before model construction. +- Per-layer dynamic EPLB requires `CUDA_SCALE_LAUNCH_QUEUES=4x` before CUDA context + initialization. TensorRT-LLM applies the setting before `TorchLlmArgs` + validation and CUDA initialization/device probes, and propagates it to MPI, + direct, and Ray workers. Applications that initialize CUDA before TensorRT-LLM + must set it at process startup. +- HALO-Q and the in-switch TMA copy run on the higher-priority copy stream. + The main stream may enqueue shared-expert work and input quantization before + it waits for the plan at the first routed-expert use. A single CPU submitter + owns both streams; the path does not require a CUDA green context or a + steady-state host wait. +- HALO-Q and TMA follow the native-operator path under + `cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb` and + `cpp/tensorrt_llm/thop/moe/loadBalance`. The scheduler accepts the original + contiguous CUDA `int32` route tensor; do not restore route staging, tail + fills, or a host-generated copy plan. +- Rebalance resources are model-scoped. Worker shutdown and construction + rollback must release the scheduler, TMA bindings, shared-slot mappings, + and fabric allocations collectively. Forward execution must not retain + activation tensors after their consumer has been enqueued. +- Per-layer dynamic EPLB selects the fused shared FC12 kernel only when the routed + backend declares per-layer dynamic support and the final projection + quantization, hardware, and graph mode support that kernel. Other runs keep + the model's existing shared-MLP backend. The active configuration supplies + the optional SM budget only to the fused kernel. ## Canonical Examples @@ -982,7 +1012,7 @@ their real class name. - **Do NOT set a communication kernel's timeout directly or through environment variables at runtime** — A backend whose kernels wait for peers under a timeout registers a `MoECommTimeoutProxy` with the timeout guard when it is constructed and unregisters it before releasing what the proxy controls; only the guard decides the value (see [Communication timeouts](#communication-timeouts)). Writing environment variables at runtime races with NCCL threads that read them - **Do NOT allocate symmetric memory from `run_moe` in `FUSED_COMM` backends** — Symmetric-memory rendezvous is a build-time collective and is unsafe under PP / layer-skip or CUDA graph capture; allocate from `create_weights()` after `ConfigurableMoE` has synchronized EPLB-derived attributes. See `mega_moe/mega_moe_deepgemm.py` for the DG pattern and `mega_moe/mega_moe_cute_dsl.py:_alloc_symm_provider` for the NVSHMEM-equivalent provider. - **Do NOT add a new `FUSED_COMM` backend without a zero-token `quantize_input` regression test** — `FusedCommMoEScheduler` calls `quantize_input` for every chunk (including zero-token chunks) so each backend must return its own empty-tensor layout. -- **Do NOT use a dataclass for an autotuner tactic without a tested `__repr__` round-trip** — `AutoTuner` serializes tactic values through `json.dumps`/`json.loads` and `eval(repr(tactic))`; a plain dataclass fails the `eval(repr(...))` check. Prefer a JSON-friendly **tuple of primitives or lists of primitives** (lists are JSON-friendly; tuples round-trip via `eval(repr(...))`). See the tactic-representation comment block in `tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py` for the 10-field tactic pattern (with legacy 8-field compatibility) (mma_tiler/cluster_shape as `list[int]`, `epi_flag_batch` as a nested `(int, int)` tuple, the rest as `bool`/`int`/`str`; `_unpack_tactic` is the single source of truth for the field order). The fallback tactic is the token-aware `default_megamoe_tactic(num_tokens)` helper, selected by `Sm100MegaMoENvfp4Runner.forward(tactic=-1)`, not a separate `fallback_tactic()` method. +- **Do NOT use a dataclass for an autotuner tactic without a tested `__repr__` round-trip** — `AutoTuner` serializes tactic values through `json.dumps`/`json.loads` and `eval(repr(tactic))`; a plain dataclass fails the `eval(repr(...))` check. Prefer a JSON-friendly **tuple of primitives or lists of primitives** (lists are JSON-friendly; tuples round-trip via `eval(repr(...))`). See the tactic-representation comment block in `tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py` for the canonical 11-field tactic pattern (with legacy 10- and 8-field compatibility), including token-back READY granularity; `_unpack_tactic` is the single source of truth for field order and compatibility defaults. The fallback tactic is the token-aware `default_megamoe_tactic(num_tokens)` helper, selected by `Sm100MegaMoENvfp4Runner.forward(tactic=-1)`, not a separate `fallback_tactic()` method. - **Use `distributed_tuning_strategy=DistributedTuningStrategy.MERGE` on a multi-rank `FUSED_COMM` backend's `TuningConfig`** — Every EP rank must converge on the same compiled tactic for every chunk, otherwise the in-kernel NVLink dispatch barrier deadlocks. `PARALLEL` can profile different tactics on different ranks and is unsafe for fused collectives. Reference: `Sm100MegaMoENvfp4Runner.get_tuning_config`. ### CUTLASS eager scratch reclamation diff --git a/tensorrt_llm/_torch/moe/fused_moe/__init__.py b/tensorrt_llm/_torch/moe/fused_moe/__init__.py index 2eade6470e55..2583c7c44fc1 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/__init__.py +++ b/tensorrt_llm/_torch/moe/fused_moe/__init__.py @@ -39,6 +39,8 @@ from .interface import MoE, MoEWeightLoadingMode from .moe_load_balancer import (MoeLoadBalancer, moe_load_balancer_set_repeated_for_next_layer) +from .per_layer_dynamic_eplb import (PerLayerDynamicMoeLayerLoadBalancer, + PerLayerDynamicMoeLoadBalancer) from .quantization import FusedMoEQuantScalesFP8 # yapf: disable from .routing import (BaseMoeRoutingMethod, DeepSeekV3MoeRoutingMethod, @@ -95,6 +97,8 @@ "MoEImplBase", "MoEImplClass", "MoeLoadBalancer", + "PerLayerDynamicMoeLayerLoadBalancer", + "PerLayerDynamicMoeLoadBalancer", "MoEWeightLoadingMode", "MiniMaxM2MoeRoutingMethod", "DeepSeekV4MoeRoutingMethod", diff --git a/tensorrt_llm/_torch/moe/fused_moe/create_moe.py b/tensorrt_llm/_torch/moe/fused_moe/create_moe.py index 37601f07b41b..074e8a8f2a31 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/moe/fused_moe/create_moe.py @@ -115,9 +115,16 @@ def create_moe_backend( # backend that cannot run under a load balancer already declines in # ``can_implement``. eplb_enabled = get_moe_load_balancer() is not None + load_balancer_config = getattr(model_config, "moe_load_balancer", None) + eplb_mode = (getattr(load_balancer_config, "mode", "standard") + if eplb_enabled else "disabled") if eplb_enabled and not moe_cls.capabilities.supports_eplb: raise ValueError( f"{moe_cls.__name__} does not support the MoE load balancer.") + if (eplb_mode == "per_layer_dynamic" + and not moe_cls.capabilities.supports_per_layer_dynamic_eplb): + raise ValueError( + f"{moe_cls.__name__} does not support per-layer dynamic EPLB.") if bias and not moe_cls.capabilities.supports_expert_bias: raise ValueError(f"bias not supported in {moe_cls.__name__}.") diff --git a/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py b/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py index a5406cc14416..1b06f098cb7c 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py +++ b/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py @@ -15,9 +15,9 @@ """Contracts for declaring, selecting, and executing MoE implementations.""" import hashlib -from dataclasses import dataclass +from dataclasses import dataclass, field from enum import Enum -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple import torch @@ -47,6 +47,9 @@ class MoEStaticCapability: # ``_supports_load_balancer()``, which TrtllmGenFusedMoEBase overrides to mean # "separated routing is used". supports_eplb: bool = False + # Whether the backend consumes a per-forward helper-slot plan whose live + # weights may be published concurrently with shared-expert computation. + supports_per_layer_dynamic_eplb: bool = False # Legacy gate: the ``assert moe_cls in [...]`` bias allow-list in # ``create_moe_backend``. Per-expert FC bias from the checkpoint, added # before the activation functor runs -- not an activation constant. @@ -287,6 +290,8 @@ class MoEDeployment: cluster_size: int = 1 # True only when an EPLB load balancer is registered. eplb_enabled: bool = False + # Unified EPLB mode: disabled, standard, or per_layer_dynamic. + eplb_mode: str = "disabled" # True only for routed-expert LoRA targets. moe_lora_enabled: bool = False # False when ``moe_disable_finalize_fusion`` is set or any LoRA is @@ -615,6 +620,7 @@ def to_dict(self) -> Dict[str, object]: "use_dp": self.deployment.use_dp, "num_slots": self.deployment.num_slots, "eplb_enabled": self.deployment.eplb_enabled, + "eplb_mode": self.deployment.eplb_mode, "moe_lora_enabled": self.deployment.moe_lora_enabled, "force_dynamic_quantization": self.deployment.force_dynamic_quantization, "oss_cutlass_moe": self.deployment.env.oss_cutlass_moe, @@ -688,6 +694,48 @@ class MoECommPlan: use_deep_ep_direct_metadata: bool = False +@dataclass(frozen=True) +class MoEReplicaPlan: + """Per-forward physical replica layout produced by the EPLB layer. + + The backend consumes these facts but does not own the planner, copy stream, + weight arena, or generation lifecycle. The private wait callback preserves + the late dependency: independent input preparation remains overlapped until + the fused kernel is about to read the physical routes. + """ + + resident_experts_per_rank: int + compute_slots_per_rank: int + ready_flags: torch.Tensor + ready_generation: int + reserved_sms: int + _wait_for_routes: Callable[[], None] = field(repr=False, compare=False) + + def __post_init__(self) -> None: + resident = int(self.resident_experts_per_rank) + compute = int(self.compute_slots_per_rank) + generation = int(self.ready_generation) + if resident <= 0 or compute <= resident: + raise ValueError( + "A replica plan requires compute slots M greater than resident " + f"experts H; got H={resident}, M={compute}." + ) + if not 1 <= generation < (1 << 63): + raise ValueError( + f"A replica plan READY generation must be in [1, 2**63); got {generation}." + ) + if int(self.reserved_sms) <= 0: + raise ValueError("A replica plan must reserve at least one auxiliary SM") + + @property + def helper_slots_per_rank(self) -> int: + return int(self.compute_slots_per_rank) - int(self.resident_experts_per_rank) + + def wait_for_routes(self) -> None: + """Order the first route consumer after the plan producer on device.""" + self._wait_for_routes() + + @dataclass(frozen=True) class MoERunContext: """Everything ``run_moe`` needs that the CALLER genuinely produces. @@ -718,6 +766,8 @@ class MoERunContext: all_rank_num_tokens: Optional[List[int]] = None # produced by the comm strategy, one object per forward comm_plan: Optional[MoECommPlan] = None + # produced by per-layer dynamic EPLB; None for standard EPLB and OFF. + replica_plan: Optional[MoEReplicaPlan] = None def require_comm_plan(impl: object, ctx: MoERunContext) -> MoECommPlan: diff --git a/tensorrt_llm/_torch/moe/fused_moe/interface.py b/tensorrt_llm/_torch/moe/fused_moe/interface.py index 3c14d8a174b6..fe1b5ef9ed2d 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/interface.py +++ b/tensorrt_llm/_torch/moe/fused_moe/interface.py @@ -514,8 +514,12 @@ def _init_load_balancer( moe_load_balancer = get_moe_load_balancer() moe_load_balancer_config = model_config.moe_load_balancer - # Calculate initial expert assignments - if moe_load_balancer_config: + # Standard EPLB checkpoints carry one expert id per physical slot. + # Per-layer dynamic EPLB keeps checkpoint/source storage on the resident + # axis H=E/EP; only the live kernel planes are widened to M=P/EP. + per_layer_dynamic_eplb = (getattr(moe_load_balancer_config, "mode", + None) == "per_layer_dynamic") + if moe_load_balancer_config and not per_layer_dynamic_eplb: init_expert_size_per_partition = moe_load_balancer_config.num_local_slots self.initial_global_assignments = [ (ep_rank * self.num_experts // self.ep_size + local_slot_id) % @@ -523,7 +527,7 @@ def _init_load_balancer( for local_slot_id in range(init_expert_size_per_partition) ] else: - # Sequential mapping: expert i → slot i; covers all experts regardless of divisibility + # Sequential resident mapping covers every logical expert exactly once. self.initial_global_assignments = list(range(self.num_experts)) # Setup load balancer if available @@ -543,23 +547,44 @@ def _init_load_balancer( f"by ep_size ({self.ep_size}) so each rank holds the same " f"number of slots.") top_k = self.routing_method.experts_per_token - self.expert_size_per_partition = moe_load_balancer_config.num_local_slots + if per_layer_dynamic_eplb: + validate_capacity = getattr(moe_load_balancer_config, + "validate_expert_capacity", None) + if callable(validate_capacity): + validate_capacity(self.num_experts) + if self.num_experts % self.ep_size != 0: + raise ValueError( + f"{type(self).__name__}: per-layer dynamic EPLB requires " + f"num_experts ({self.num_experts}) divisible by ep_size " + f"({self.ep_size}) so resident H is rank-identical.") + self.expert_size_per_partition = self.num_experts // self.ep_size + else: + self.expert_size_per_partition = moe_load_balancer_config.num_local_slots - # Add this layer to the load balancer + # Add this layer to the load balancer. slot_count_per_rank is the + # compute capacity M for per-layer dynamic mode and the ordinary local + # slot count for standard mode. aux_stream = getattr(self, '_get_load_balancer_aux_stream', lambda: None)() + slot_count_per_rank = (moe_load_balancer_config.num_local_slots + if per_layer_dynamic_eplb else + self.expert_size_per_partition) self.layer_load_balancer = moe_load_balancer.add_layer( self.num_experts, top_k, - self.expert_size_per_partition, + slot_count_per_rank, aux_stream=aux_stream) self.repeat_count = self.layer_load_balancer.get_repeat_count() # Handle initial global assignments - loaded_initial_global_assignments = ( - moe_load_balancer_config.get_layer_initial_global_assignments( - self.layer_idx)) + get_initial_assignments = getattr( + moe_load_balancer_config, + "get_layer_initial_global_assignments", + None, + ) + loaded_initial_global_assignments = (get_initial_assignments( + self.layer_idx) if callable(get_initial_assignments) else None) self.num_slots = moe_load_balancer_config.num_slots if loaded_initial_global_assignments is not None: @@ -581,7 +606,8 @@ def _init_load_balancer( f"initial_global_assignments (layer {self.layer_idx}) = {self.initial_global_assignments}" ) - # Slot boundaries for EPLB (uniform: all ranks hold same num_local_slots) + # Source/checkpoint boundaries are resident H in per-layer dynamic + # mode. The layer manager separately owns compute slots M. self.slot_start = self.ep_rank * self.expert_size_per_partition self.slot_end = self.slot_start + self.expert_size_per_partition self.initial_local_expert_ids = self.initial_global_assignments[ @@ -675,13 +701,41 @@ def _load_balancer_update_statistic(self, def _load_balancer_route(self, token_selected_experts: torch.Tensor, use_dp: bool) -> torch.Tensor: - """Route tokens using load balancer.""" - if self.layer_load_balancer: - return self.layer_load_balancer.route(token_selected_experts, - use_dp) - else: + """Route tokens and enqueue a per-layer dynamic replica plan when enabled.""" + if not self.layer_load_balancer: return token_selected_experts + gap_hook = getattr(self, "_rebalance_plan_gap_hook", None) + if gap_hook is not None: + # The hook is one-shot: DeepSeekV4 falls back to ordinary shared + # expert execution if routing exits before consuming it. + self._rebalance_plan_gap_hook = None + route = self.layer_load_balancer.route + if gap_hook is not None: + return route(token_selected_experts, use_dp, gap_hook=gap_hook) + return route(token_selected_experts, use_dp) + + def _load_balancer_uses_replica_plan(self) -> bool: + """Whether routing is produced by the per-layer dynamic device planner.""" + return bool(self.layer_load_balancer and callable( + getattr(self.layer_load_balancer, "require_replica_plan", None))) + + def _load_balancer_replica_plan(self): + """Return the plan produced by the most recent route, if this mode has one.""" + if not self._load_balancer_uses_replica_plan(): + return None + return self.layer_load_balancer.require_replica_plan() + + def _load_balancer_finish_replica_plan(self, plan) -> None: + """Release a caller-owned plan after its last MAIN-stream consumer.""" + if plan is None or not self.layer_load_balancer: + return + finish = getattr(self.layer_load_balancer, "finish_replica_plan", None) + if not callable(finish): + raise RuntimeError( + "an EPLB replica plan was produced without a finish method") + finish(plan) + def _load_balancer_start_set_cpu_stage(self, is_last_call: bool): """Start CPU stage in load balancer.""" if self._using_dynamic_load_balancer() and is_last_call: diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py index 638892f0f871..77f19f920bef 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py @@ -70,7 +70,6 @@ import os import socket -import time import weakref from dataclasses import dataclass from typing import Dict, List, Optional, Tuple, Union @@ -141,6 +140,8 @@ class MegaMoeCuteDslUnavailable(RuntimeError): _MEGAMOE_PRIME_LADDER = os.environ.get("MEGAMOE_PRIME_LADDER", "1") == "1" _MEGAMOE_PRIMED_LADDERS: set = set() +# Distinguish an absent arm attribute from an explicit False when restoring it. + def is_megamoe_cute_dsl_runtime_available() -> Tuple[bool, Optional[str]]: """Return whether the CUDA 13 Cutlass DSL runtime exposes all symbols the @@ -240,7 +241,7 @@ def is_megamoe_cute_dsl_runtime_available() -> Tuple[bool, Optional[str]]: return False, reason try: - from ....moe.kernels.cutedsl_megamoe import ( # noqa: F401 + from ....moe.kernels.megamoe_kernels import ( # noqa: F401 BlackwellInferenceMegaMoE, RubinInferenceGenphaseMegaMoE, RubinInferenceLocalMegaMoE, @@ -278,6 +279,7 @@ def is_megamoe_cute_dsl_runtime_available() -> Tuple[bool, Optional[str]]: # ``_to_cute``. +# The launch path memoizes these dtype views without copying their storage. def _as_nvfp4(t: torch.Tensor) -> torch.Tensor: return t if t.dtype == torch.float4_e2m1fn_x2 else t.view(torch.float4_e2m1fn_x2) @@ -404,7 +406,10 @@ class TrtllmCutedslMegaMoeNvfp4Impl(MoEImplBase): # Static and dynamic EPLB both work: see ``_supports_load_balancer`` # below for why the MegaMoE-format derived parameters migrate # atomically. - capabilities=MoEStaticCapability(supports_eplb=True), + capabilities=MoEStaticCapability( + supports_eplb=True, + supports_per_layer_dynamic_eplb=True, + ), doc="MegaMoE CuteDSL fused NVFP4 kernels: NVFP4 weights and activations, SM100/103/107.", ) @@ -430,6 +435,16 @@ class TrtllmCutedslMegaMoeNvfp4Impl(MoEImplBase): # ``CombineFormat.parse`` in the kernel package's token_comm.py. _SUPPORTED_COMBINE_FORMATS = frozenset({"bf16", "32e4m3xe8m0", "16e2m1xbf16"}) + # Disabled defaults also cover construction paths before weight allocation. + _rebalance_slots_active: int = 0 + _quantize_reserved_sms: int = 0 + # Resident count H is recorded at weight creation and never widened. + _rebalance_home_experts: Optional[int] = None + # Active per-layer dynamic variant of ``model_config.moe_load_balancer``. + _per_layer_dynamic_eplb_config = None + # Allocate the arena once; the copy endpoint retains its plane identities. + _rebalance_arena = None + # ------------------------------------------------------------------ # Capability gating # ------------------------------------------------------------------ @@ -542,6 +557,23 @@ def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: # ------------------------------------------------------------------ # Init # ------------------------------------------------------------------ + @staticmethod + def _reject_rebalance_dwdp_overlap(mapping, helper_slots: int, layer_idx=None) -> None: + """Reject simultaneous helper-slot rebalance and DWDP before allocation. + + DWDP makes every rank own all experts and rewrites the slot axis after + weight creation, invalidating the H + S arena geometry. + """ + if int(helper_slots) <= 0: + return + if not bool(getattr(mapping, "dwdp_enabled", False)): + return + raise ValueError( + f"MoE rebalance S={helper_slots} at layer {layer_idx} cannot run with " + f"DWDP size {getattr(mapping, 'dwdp_size', '?')}: DWDP rewrites the " + "H + S weight axis. Disable rebalance or use dwdp_size=1." + ) + def __init__( self, *, @@ -598,10 +630,7 @@ def __init__( ) self._topk_idx_dtype = torch.int32 if get_sm_version() == 107 else torch.int64 - # Cross-rank combine wire format is selected with - # MEGAMOE_COMBINE_FORMAT (default bf16). It MUST be rank-identical: - # it sizes the symmetric workspace and picks the compiled kernel, so - # divergence desyncs the rendezvous / NVLink barrier. + # The wire format must match on all EP ranks: it sizes the symmetric workspace. combine_format = os.environ.get("MEGAMOE_COMBINE_FORMAT", "bf16") if combine_format not in self._SUPPORTED_COMBINE_FORMATS: raise ValueError( @@ -616,6 +645,62 @@ def __init__( # gates collective tuning behavior). self.tactic_autotune = os.environ.get("MEGAMOE_TACTIC_AUTOTUNE", "0") == "1" + # Per-layer dynamic EPLB uses the standard load-balancer configuration. + # P is the global physical slot count; the fused kernel sees + # M = P / EP slots per rank and S = M - H helper rows. + load_balancer = getattr(model_config, "moe_load_balancer", None) + self._per_layer_dynamic_eplb_config = ( + load_balancer if getattr(load_balancer, "mode", None) == "per_layer_dynamic" else None + ) + self._rebalance_autotune_helpers_initialized = False + self._rebalance_slots_active = 0 + if self._per_layer_dynamic_eplb_config is not None: + ep_size = int(self.ep_size) + global_slots = int(self._per_layer_dynamic_eplb_config.num_slots) + if ep_size <= 0 or num_experts <= 0 or num_experts % ep_size != 0: + raise ValueError( + "per-layer dynamic EPLB requires a positive expert count divisible " + f"by EP; got E={num_experts}, EP={ep_size}." + ) + if global_slots % ep_size != 0: + raise ValueError( + "per-layer dynamic EPLB requires num_slots divisible by EP; " + f"got P={global_slots}, EP={ep_size}." + ) + resident_experts = num_experts // ep_size + compute_slots = global_slots // ep_size + self._rebalance_slots_active = compute_slots - resident_experts + if self._rebalance_slots_active <= 0: + raise ValueError( + "per-layer dynamic EPLB requires physical capacity larger than " + f"resident capacity; got P={global_slots}, E={num_experts}." + ) + + self._quantize_reserved_sms = 0 + self._rebalance_reserved_sms = 0 + if self._rebalance_slots_active > 0: + self._rebalance_reserved_sms = int(self._per_layer_dynamic_eplb_config.auxiliary_sms) + if hasattr(torch.ops.trtllm.fp4_quantize, "sm_budget"): + self._quantize_reserved_sms = self._rebalance_reserved_sms + else: + logger.warning_once( + "fp4_quantize.sm_budget is unavailable; DynamicEPLB " + "quantization will use the selected backend without SM control", + key="dynamic_eplb_quantize_without_sm_budget", + ) + # Reject DWDP before it can rewrite the same slot axis. + self._reject_rebalance_dwdp_overlap( + getattr(self, "mapping", None), + self._rebalance_slots_active, + layer_idx, + ) + if self._rebalance_slots_active > 0: + logger.debug( + f"[MegaMoECuteDsl] prefill rebalance active " + f"(S={self._rebalance_slots_active}); kernel weight planes " + f"use M = H + S while expert_size_per_partition stays H." + ) + # Buffer sizing. MoE layers execute serially per forward; one pool # sized to the worst-case per-rank tokens covers every layer. The # kernel compile takes this as the static ``max_tokens_per_rank``. @@ -689,9 +774,33 @@ def __init__( # scalar would reset the wrong buffer when buckets alternate). self._local_staging_cache: Dict[Tuple, Dict[str, torch.Tensor]] = {} self._last_staged_T: Dict[int, int] = {} + # Cache immutable provider regions and dtype views, retaining their owners. + # Hits validate object identity; dtype views also detect storage rebinding. + self._regions_memo: Dict[int, Tuple] = {} + self._dtype_view_memo: Dict[Tuple, Tuple] = {} if not model_config.skip_create_weights_in_init: self.create_weights() + # ------------------------------------------------------------------ + # Compute axis + # ------------------------------------------------------------------ + @property + def compute_slot_count(self) -> int: + """Return the kernel-only slot axis ``M = H + S``. + + ``expert_size_per_partition`` remains the resident/checkpoint axis H. + Only the seven live kernel planes, symmetric workspace, and launch + geometry use this wider axis. + """ + resident_experts = int(self.expert_size_per_partition) + recorded_home = self._rebalance_home_experts + if recorded_home is not None and resident_experts != recorded_home: + raise RuntimeError( + "MegaMoECuteDsl resident expert geometry changed after weight " + f"creation: expected H={recorded_home}, got H={resident_experts}." + ) + return resident_experts + int(self._rebalance_slots_active) + # ------------------------------------------------------------------ # Topology # ------------------------------------------------------------------ @@ -944,12 +1053,15 @@ def create_weights(self): """Build-time weight + symmetric-buffer allocation. Order: + 0. Record the resident/checkpoint axis H. The backend keeps + ``expert_size_per_partition == H``; ``compute_slot_count`` exposes + ``M = H + S`` only to kernel-plane allocation and launch geometry. 1. Allocate symmetric-memory provider for multi-rank EP (collective rendezvous; MUST run at build time -- not from ``run_moe`` -- because forward time may be inside CUDA graph capture or non-lockstep PP/layer-skip). - 2. Resolve quantization method. - 3. Delegate parameter registration to the quant method. + 2. Resolve the quantization method and its seven-plane ABI. + 3. Delegate parameter registration to the same quant method. 4. Flip ``_weights_created``. The symm provider is shared across MoE layers with the same @@ -961,13 +1073,67 @@ def create_weights(self): """ if self._weights_created: return + # ConfigurableMoE has synchronized resident geometry H before this call. + # Keep that public/checkpoint axis unchanged. The explicit + # ``compute_slot_count`` property is the only H + S axis. + resident_experts = int(self.expert_size_per_partition) + helper_slots = int(self._rebalance_slots_active) + if self._rebalance_home_experts is None: + self._rebalance_home_experts = resident_experts + elif resident_experts != self._rebalance_home_experts: + raise RuntimeError( + "MegaMoECuteDsl: expert_size_per_partition changed between " + f"create_weights() calls -- expected resident H=" + f"{self._rebalance_home_experts}, found H={resident_experts}." + ) + if helper_slots > 0: + log_fn = logger.info if self.layer_idx in (0, None) else logger.debug + log_fn( + f"[MegaMoECuteDsl] layer={self.layer_idx} MoE rebalance uses " + f"resident H={resident_experts} and helper S={helper_slots}; " + f"only the seven live kernel planes use " + f"M={self.compute_slot_count} rows." + ) + # The quant method is the sole authority for the seven live planes. + if self.quant_method is None: + self.quant_method = self._get_quant_method() + live_weight_plane_spec = self.quant_method.live_weight_plane_spec(self) + # Step 1: build-time symmetric memory allocation (multi-rank only). # Single-rank degenerate uses local CUDA tensors and skips here. if self.ep_size > 1: self._alloc_symm_provider() - # Step 2-3: quant method registers all NVFP4 + MegaMoE-format params. - self.quant_method = self._get_quant_method() + # Allocate the H + S arena before the quant method creates the seven + # kernel-plane Parameters. Source/checkpoint Parameters remain H rows. + if self._rebalance_slots_active > 0 and getattr(self, "_rebalance_arena", None) is None: + # The hierarchical multicast arena: one multicast group per HALO + # level, rather than one per (source, destination) pair. + from .rebalance_live_arena import allocate_rebalance_arena + + layer_balancer = getattr(self, "layer_load_balancer", None) + manager = getattr(layer_balancer, "manager", None) + create_provider = getattr(manager, "create_arena_provider", None) + if not callable(create_provider): + raise RuntimeError("per-layer dynamic EPLB requires a model-owned arena provider") + get_ep_comm = getattr(manager, "get_ep_comm", None) + if not callable(get_ep_comm): + raise RuntimeError("per-layer dynamic EPLB requires a model-owned EP communicator") + self._rebalance_arena = allocate_rebalance_arena( + home_experts=int(self._rebalance_home_experts), + helper_slots=int(self._rebalance_slots_active), + bundle=live_weight_plane_spec, + ep_comm=get_ep_comm(self._ep_pg), + mapping=self.mapping, + provider=create_provider(), + device=int(self.mapping.local_rank), + layer_idx=self.layer_idx, + ) + # Step 2-3: the same quant method registers every live Parameter. self.quant_method.create_weights(self) + if self._rebalance_arena is not None: + from .rebalance_slot_scheduler import bind_live_weight_planes + + bind_live_weight_planes(self.quant_method, self) # Step 4. self._weights_created = True @@ -1018,7 +1184,7 @@ def _alloc_symm_provider(self): ) = query_megamoe_kernel_properties( world_size=self.ep_size, num_topk=top_k, - num_experts_per_rank=int(self.expert_size_per_partition), + num_experts_per_rank=self.compute_slot_count, hidden_size=self.hidden_size, intermediate_size_per_partition=int(self.intermediate_size_per_partition), expand_intermediate_size_per_partition=int( @@ -1125,6 +1291,17 @@ def cache_derived_state(self) -> None: if self.quant_method is None: self.create_weights() super().cache_derived_state() + # ConfigurableMoE delegates this hook on both ordinary and staged + # loading paths. The layer manager builds and owns the producer only + # after the final seven plane aliases exist. + if self._rebalance_slots_active > 0: + arena = self._rebalance_arena + arena.assert_identity() + self.quant_method.assert_live_weight_aliases(self, arena) + bind_backend = getattr(self.layer_load_balancer, "bind_backend", None) + if not callable(bind_backend): + raise RuntimeError("per-layer dynamic EPLB layer manager cannot bind its runtime") + bind_backend(self) def pre_reload_weights(self) -> None: raise NotImplementedError( @@ -1138,7 +1315,7 @@ def _build_weight_view(self) -> MegaMoECuteDslWeightView: kernel sees the latest dynamic-EPLB migration outcome (once that path lands; currently the slots are static). """ - return MegaMoECuteDslWeightView( + _view = MegaMoECuteDslWeightView( fc1_weight=self.mega_fc1_weight, fc1_weight_sf=self.mega_fc1_weight_sf, fc2_weight=self.mega_fc2_weight, @@ -1147,6 +1324,7 @@ def _build_weight_view(self) -> MegaMoECuteDslWeightView: fc2_alpha=self.fc2_alpha, fc1_norm_const=self.fc1_norm_const, ) + return _view # ------------------------------------------------------------------ # MoE-contract methods @@ -1200,13 +1378,25 @@ def quantize_input( # admits bf16 activations only, so ``fc31_act_scale``, created at # ``module.dtype``, is bf16 as well. x_bf16 = x_bf16 * fc31_act_scale - x_fp4, x_sf = torch.ops.trtllm.fp4_quantize( - x_bf16, - self.fc31_input_scale, - 16, # scaling_vector_size == Nvfp4BlockSize - False, # sf_use_ue8m0 - False, # is_sf_swizzled - MegaMoE expects plain K-major - ) + # Match the concurrent scheduler/copy reservation. The default overload + # remains unchanged when load balancing is disabled. + if self._quantize_reserved_sms: + x_fp4, x_sf = torch.ops.trtllm.fp4_quantize.sm_budget( + x_bf16, + self.fc31_input_scale, + 16, + False, + False, + reserved_sms=self._quantize_reserved_sms, + ) + else: + x_fp4, x_sf = torch.ops.trtllm.fp4_quantize( + x_bf16, + self.fc31_input_scale, + 16, # scaling_vector_size == Nvfp4BlockSize + False, # sf_use_ue8m0 + False, # is_sf_swizzled - MegaMoE expects plain K-major + ) # ``fp4_quantize(is_sf_swizzled=False)`` returns LINEAR layout # ``(rows, ceil(hidden/16))`` with no column pad. The kernel TMA # load needs ``pad_up(ceil_div(hidden, 16), 4)`` bytes per row @@ -1221,6 +1411,49 @@ def quantize_input( padded_sf[:, :raw_cols] = x_sf_raw return x_fp4, padded_sf + def _replica_plan_launch_args(self, replica_plan): + """Validate and unpack the caller-owned per-forward replica plan.""" + if replica_plan is None: + return None + resident_experts = int(replica_plan.resident_experts_per_rank) + compute_slots = int(replica_plan.compute_slots_per_rank) + helper_slots = compute_slots - resident_experts + expected_resident = int(self.expert_size_per_partition) + expected_compute = self.compute_slot_count + if resident_experts != expected_resident or compute_slots != expected_compute: + raise RuntimeError( + "MegaMoECuteDsl replica-plan geometry does not match its allocated " + f"weights: plan H/M={resident_experts}/{compute_slots}, " + f"backend H/M={expected_resident}/{expected_compute}." + ) + generation = int(replica_plan.ready_generation) + reserved_sms = int(replica_plan.reserved_sms) + if helper_slots <= 0 or not 1 <= generation < (1 << 63) or reserved_sms <= 0: + raise RuntimeError( + "MegaMoECuteDsl received an invalid replica plan: " + f"S={helper_slots}, generation={generation}, " + f"reserved_sms={reserved_sms}." + ) + quantize_reserved_sms = int(getattr(self, "_quantize_reserved_sms", 0)) + if quantize_reserved_sms not in (0, reserved_sms): + raise RuntimeError( + "MegaMoECuteDsl replica-plan SM reservation differs from the " + f"quantization reservation: plan={reserved_sms}, " + f"quantize={quantize_reserved_sms}." + ) + return ( + compute_slots, + helper_slots, + replica_plan.ready_flags, + generation, + reserved_sms, + ) + + def _wait_rebalance_routes(self, replica_plan=None) -> None: + """Order route consumers after HALO-Q while leaving setup overlapped.""" + if replica_plan is not None: + replica_plan.wait_for_routes() + def run_moe( self, ctx: MoERunContext, @@ -1239,6 +1472,7 @@ def run_moe( token_selected_experts = ctx.token_selected_experts token_final_scales = ctx.token_final_scales x_sf = ctx.x_sf + replica_plan = getattr(ctx, "replica_plan", None) output_dtype = ctx.output_dtype if output_dtype is None: output_dtype = self.dtype or torch.bfloat16 @@ -1271,8 +1505,15 @@ def run_moe( top_k = int(token_selected_experts.shape[-1]) device = x.device - topk_idx = token_selected_experts.to(self._topk_idx_dtype).contiguous() topk_weights_f32 = token_final_scales.to(torch.float32).contiguous() + # The direct scheduler path borrows contiguous int32 output without a GPU + # operation. Any required conversion is itself a route reader. + if ( + token_selected_experts.dtype != self._topk_idx_dtype + or not token_selected_experts.is_contiguous() + ): + self._wait_rebalance_routes(replica_plan) + topk_idx = token_selected_experts.to(self._topk_idx_dtype).contiguous() return self._run_moe( x=x, @@ -1285,6 +1526,7 @@ def run_moe( hidden=hidden, device=device, output_dtype=output_dtype, + replica_plan=replica_plan, ) def _ensure_local_staging(self, *, top_k: int, hidden: int, device, output_dtype, max_T: int): @@ -1355,7 +1597,7 @@ def _ensure_local_staging(self, *, top_k: int, hidden: int, device, output_dtype shared_bytes, supports_direct_inputs = query_megamoe_kernel_properties( world_size=1, num_topk=top_k, - num_experts_per_rank=int(self.expert_size_per_partition), + num_experts_per_rank=self.compute_slot_count, hidden_size=hidden, intermediate_size_per_partition=int(self.intermediate_size_per_partition), expand_intermediate_size_per_partition=int( @@ -1388,6 +1630,33 @@ def _ensure_local_staging(self, *, top_k: int, hidden: int, device, output_dtype cached[cache_key] = staging return staging + def _memo_regions(self, provider): + """Cache views of an immutable provider, retaining its identity and lifetime.""" + key = id(provider) + hit = self._regions_memo.get(key) + # Retain the provider and check identity to prevent object-ID reuse. + if hit is not None and hit[0] is provider: + return hit[1] + regions = provider.get_regions() + self._regions_memo[key] = (provider, regions) + return regions + + def _memo_dtype_view(self, t: torch.Tensor, caster) -> torch.Tensor: + """Cache dtype views while retaining the source tensor. + + In-place helper-weight copies remain visible through these aliases. + Storage rebinding must invalidate the cached view even if identity matches. + """ + key = (id(t), caster) + hit = self._dtype_view_memo.get(key) + # Check identity and data_ptr: assigning Parameter.data can replace storage + # without changing the object. These dtype views preserve shape and stride. + if hit is not None and hit[0] is t and hit[1].data_ptr() == t.data_ptr(): + return hit[1] + v = caster(t) + self._dtype_view_memo[key] = (t, v) + return v + def _acquire_buffers( self, *, top_k: int, hidden: int, device, output_dtype, launch_max_T: int ) -> _MegaMoeBuffers: @@ -1439,7 +1708,7 @@ def _acquire_buffers( f"{provider.num_topk} but run_moe called with " f"top_k={top_k}; recreate the backend." ) - regions = provider.get_regions() + regions = self._memo_regions(provider) return _MegaMoeBuffers( activation=regions.activation, activation_sf=regions.activation_sf, @@ -1461,6 +1730,8 @@ def _stage_inputs( num_tokens: int, top_k: int, stage_activation: bool, + stage_routes: bool = True, + replica_plan=None, ) -> None: """Copy live rows of the user-domain inputs into the kernel's pre-allocated buffers and refresh the padded tail. @@ -1487,17 +1758,20 @@ def _stage_inputs( """ max_T = bufs.topk_idx_local.shape[0] last_T = self._last_staged_T.get(max_T) - if last_T is not None and last_T > num_tokens: + if stage_routes and last_T is not None and last_T > num_tokens: bufs.topk_idx_local[num_tokens:last_T].fill_(-1) if num_tokens > 0: - bufs.topk_idx_local[:num_tokens].copy_(topk_idx, non_blocking=True) if stage_activation: bufs.activation[:num_tokens].copy_(x.view(torch.uint8), non_blocking=True) bufs.activation_sf[:num_tokens].copy_(x_sf.view(torch.uint8), non_blocking=True) bufs.topk_weights[:num_tokens, :top_k].copy_(topk_weights, non_blocking=True) if num_tokens < max_T: bufs.topk_weights[num_tokens:max_T, :top_k].zero_() - self._last_staged_T[max_T] = num_tokens + if stage_routes: + if num_tokens > 0: + self._wait_rebalance_routes(replica_plan) + bufs.topk_idx_local[:num_tokens].copy_(topk_idx, non_blocking=True) + self._last_staged_T[max_T] = num_tokens def _launch_megamoe_kernel( self, @@ -1517,6 +1791,7 @@ def _launch_megamoe_kernel( num_tokens: int, output_dtype: torch.dtype, launch_max_T: int, + replica_plan=None, ) -> torch.Tensor: """Launch the fused MegaMoE CuteDSL kernel and return its reduced output. @@ -1568,22 +1843,48 @@ def _deferred_scratch_factory(_self=self, _get=get_megamoe_profiling_scratch): if in_kernel_fc2_reduce and num_tokens > 0: combine_output[:num_tokens].zero_() + # Do not memoize scratch-provider regions: a strong reference would keep + # its multi-GiB allocation alive after profiling scratch is released. set_active_megamoe_profiling_scratch( scratch_provider.get_regions() if (scratch_provider is not None and world_size > 1) else None ) set_active_megamoe_profiling_scratch_factory(scratch_factory) + # READY generation comes only from the caller-owned replica plan. + # Helper-free launches (OFF and internal bucket priming) keep the + # M-shaped allocation but compile out the READY gate with S=0. + helper_ready_flags = None + helper_ready_generation = 0 + compute_slots = self.compute_slot_count + plan_args = self._replica_plan_launch_args(replica_plan) + if plan_args is None: + _S_active = 0 + else: + ( + compute_slots, + _S_active, + helper_ready_flags, + helper_ready_generation, + _reserved_sms, + ) = plan_args + # Record the actual helper-slot ABI for verification. + self._last_S_to_kernel = _S_active + # Quantization, independent staging, output clears, and host setup + # are already enqueued. The router inside this op first reads routes. + self._wait_rebalance_routes(replica_plan) try: torch.ops.trtllm.cute_dsl_megamoe_nvfp4_blackwell( + # Activations are per-forward allocations; only persistent weights + # belong in the strong-reference dtype-view memo. activation=_as_nvfp4(activation), activation_sf=_as_fp8_sf(activation_sf), topk_idx=topk_idx, topk_weights=topk_weights, - fc1_weight=_as_nvfp4(weight_view.fc1_weight), - fc1_weight_sf=_as_fp8_sf(weight_view.fc1_weight_sf), - fc2_weight=_as_nvfp4(weight_view.fc2_weight), - fc2_weight_sf=_as_fp8_sf(weight_view.fc2_weight_sf), + fc1_weight=self._memo_dtype_view(weight_view.fc1_weight, _as_nvfp4), + fc1_weight_sf=self._memo_dtype_view(weight_view.fc1_weight_sf, _as_fp8_sf), + fc2_weight=self._memo_dtype_view(weight_view.fc2_weight, _as_nvfp4), + fc2_weight_sf=self._memo_dtype_view(weight_view.fc2_weight_sf, _as_fp8_sf), fc1_alpha=weight_view.fc31_alpha, fc2_alpha=weight_view.fc2_alpha, fc1_norm_const=weight_view.fc1_norm_const, @@ -1592,7 +1893,7 @@ def _deferred_scratch_factory(_self=self, _get=get_megamoe_profiling_scratch): world_size=world_size, local_rank=local_rank, num_topk=top_k, - num_experts_per_rank=int(self.expert_size_per_partition), + num_experts_per_rank=compute_slots, hidden_size=hidden, intermediate_size_per_partition=int(self.intermediate_size_per_partition), expand_intermediate_size_per_partition=int( @@ -1615,6 +1916,10 @@ def _deferred_scratch_factory(_self=self, _get=get_megamoe_profiling_scratch): combine_format=self.combine_format, tactic_autotune=bool(self.tactic_autotune), num_tokens=num_tokens, + # Helper-free launches omit the READY gate. + helper_expert_count=_S_active, + hot_expert_weight_ready_flags=helper_ready_flags, + hot_expert_weight_ready_generation=int(helper_ready_generation), ) finally: # Always clear: a raising op call must not leave the factory @@ -1651,7 +1956,7 @@ def _prime_ladder( self.ep_size, hidden, int(self.intermediate_size_per_partition), - int(self.expert_size_per_partition), + self.compute_slot_count, top_k, self.combine_format, self.act_clamp, @@ -1663,8 +1968,35 @@ def _prime_ladder( ) if key in _MEGAMOE_PRIMED_LADDERS: return - t0 = time.perf_counter() primed = [] + self._prime_ladder_buckets( + live_T=live_T, + top_k=top_k, + hidden=hidden, + device=device, + output_dtype=output_dtype, + weight_view=weight_view, + primed=primed, + ) + _MEGAMOE_PRIMED_LADDERS.add(key) + if self.ep_rank == 0: + logger.debug( + "[MegaMoECuteDsl] adaptive bucket priming complete: count=%s", + len(primed), + ) + + def _prime_ladder_buckets( + self, + *, + live_T, + top_k, + hidden, + device, + output_dtype, + weight_view, + primed, + ): + """Prime each max-token bucket; the caller manages the rebalance arm state.""" for max_T in self._maxt_buckets: if max_T == live_T: continue @@ -1685,6 +2017,7 @@ def _prime_ladder( top_k=top_k, stage_activation=False, ) + # Bucket priming has no replica plan and therefore compiles the S=0 path. self._launch_megamoe_kernel( activation=bufs.activation, activation_sf=bufs.activation_sf, @@ -1703,12 +2036,6 @@ def _prime_ladder( launch_max_T=max_T, ) primed.append(max_T) - _MEGAMOE_PRIMED_LADDERS.add(key) - if self.ep_rank == 0: - logger.info( - f"[MegaMoECuteDsl] primed adaptive buckets {primed} (live bucket " - f"{live_T}) in {time.perf_counter() - t0:.1f}s" - ) def _run_moe( self, @@ -1723,6 +2050,7 @@ def _run_moe( hidden: int, device, output_dtype: torch.dtype, + replica_plan=None, ) -> torch.Tensor: """Unified MegaMoE CuteDSL forward: acquire -> stage -> launch. @@ -1802,6 +2130,7 @@ def _run_moe( and topk_idx.data_ptr() % 16 == 0 and topk_weights.data_ptr() % 16 == 0 ) + borrowed_routes = None if not use_direct_routing_inputs: self._stage_inputs( bufs=bufs, @@ -1812,11 +2141,19 @@ def _run_moe( num_tokens=num_tokens, top_k=top_k, stage_activation=not use_direct_inputs, + stage_routes=borrowed_routes is None, + replica_plan=replica_plan, ) return self._launch_megamoe_kernel( activation=x if use_direct_inputs else bufs.activation, activation_sf=x_sf if use_direct_inputs else bufs.activation_sf, - topk_idx=topk_idx if use_direct_routing_inputs else bufs.topk_idx_local, + topk_idx=( + borrowed_routes + if borrowed_routes is not None + else topk_idx + if use_direct_routing_inputs + else bufs.topk_idx_local + ), topk_weights=( topk_weights if use_direct_routing_inputs else bufs.topk_weights[:, :top_k] ), @@ -1831,6 +2168,7 @@ def _run_moe( num_tokens=num_tokens, output_dtype=output_dtype, launch_max_T=launch_max_T, + replica_plan=replica_plan, ) diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py index 37f50886ee56..224153392bc9 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py @@ -188,28 +188,34 @@ def release_symm_buffer_cache() -> None: """ if not _MEGA_MOE_SYMM_BUFFER_CACHE: return - # Detach before freeing so a raising free cannot leave a half-destroyed - # buffer in the cache for a later lookup to trip over. - entries = list(_MEGA_MOE_SYMM_BUFFER_CACHE.values()) + # Detach before freeing so successful entries cannot be reused while + # teardown is in progress. Failed entries are restored below so the worker + # can retry them before destroying their EP process group. + entries = list(_MEGA_MOE_SYMM_BUFFER_CACHE.items()) _MEGA_MOE_SYMM_BUFFER_CACHE.clear() - # Free one at a time and keep going on failure: the entries are already - # out of the cache, so any buffer skipped by a raising predecessor would - # be unreclaimable for the rest of the process lifetime. released = 0 total_bytes = 0 - for buffered, _ in entries: + first_error = None + for key, (buffered, process_group) in entries: try: total_bytes += _free_symm_buffer(buffered) released += 1 - except Exception as e: + except Exception as error: + _MEGA_MOE_SYMM_BUFFER_CACHE[key] = (buffered, process_group) + if first_error is None: + first_error = error logger.error( f"[MegaMoE] failed to release a DG SymmBuffer during executor " - f"teardown, continuing with the remaining buffers: {e}" + f"teardown, retaining it for retry: {error}" ) logger.info( f"[MegaMoE] released {released}/{len(entries)} DG SymmBuffer(s): " f"{total_bytes / 2**30:.2f} GiB" ) + if first_error is not None: + raise RuntimeError( + f"failed to release {len(_MEGA_MOE_SYMM_BUFFER_CACHE)} MegaMoE DG SymmBuffer(s)" + ) from first_error # ---- Fused MXFP8 per-token quant backends -------------------------------- diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_live_arena.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_live_arena.py new file mode 100644 index 000000000000..0d4566f3af44 --- /dev/null +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_live_arena.py @@ -0,0 +1,1058 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""EP-scoped live-weight storage for HALO-Q and hierarchical SAMI copies. + +Every layer owns its home-expert storage and maps one of two shared helper-slot +banks directly after it in virtual address space. Layers alternate banks, while +per-layer READY terminals and small planes remain independent. Allocation and +binding are collective across the EP group; vendor imports stay lazy until the +rebalance path is enabled. +""" + +import math +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple + +import torch +from torch.utils._python_dispatch import _disable_current_modes + +from tensorrt_llm.logger import logger + +__all__ = ["SharedSlotArenaProvider", "allocate_rebalance_arena"] + +BUFFER_ALIGNMENT = 2 * 1024 * 1024 + +#: The ``uint64[EP]`` terminal array is addressed by the SAMI copy kernel and +#: wants its own page. +TERMINAL_ALIGNMENT = 4096 + + +def _align_up(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +def _plane_dtypes() -> Dict[str, torch.dtype]: + """Return the canonical kernel-plane dtypes. + + Resolve optional Torch dtypes lazily so unsupported wheels fail only when + constructing an enabled rebalance arena. + """ + fp4 = getattr(torch, "float4_e2m1fn_x2", None) + fp8 = getattr(torch, "float8_e4m3fn", None) + if fp4 is None or fp8 is None: + raise RuntimeError( + "MoE rebalance live arena needs torch.float4_e2m1fn_x2 and " + "torch.float8_e4m3fn to type the NVFP4 weight planes; this torch " + f"({torch.__version__}) exposes float4_e2m1fn_x2=" + f"{fp4 is not None}, float8_e4m3fn={fp8 is not None}." + ) + return { + "uint8": torch.uint8, + "float4_e2m1fn_x2": fp4, + "float8_e4m3fn": fp8, + "float32": torch.float32, + } + + +def _plane_view( + flat_u8: torch.Tensor, + plane: Any, + slots: int, + *, + kernel: bool, +) -> torch.Tensor: + """Create a kernel or framework view from the canonical plane descriptor.""" + metadata = plane.kernel_metadata(slots) if kernel else plane.storage_metadata(slots) + shape, stride, dtype_name, element_size, pointer_alignment = metadata + dtype = _plane_dtypes().get(dtype_name) + if dtype is None: + raise RuntimeError(f"unsupported live-plane dtype {dtype_name!r}") + typed = flat_u8.view(dtype) + if typed.element_size() != element_size: + raise RuntimeError( + f"{plane.name} descriptor element size {element_size} does not " + f"match torch dtype {dtype_name} ({typed.element_size()})" + ) + if int(typed.data_ptr()) % pointer_alignment: + raise RuntimeError(f"{plane.name} pointer is not {pointer_alignment}-byte aligned") + return typed.as_strided(shape, stride) + + +@dataclass +class RebalanceLiveArena: + """Own a bound live arena and both views of its seven weight planes. + + Kernel views preserve the allocator's bound identities. Framework aliases + back nn.Parameters without changing the allocation or plane offsets. + """ + + #: The allocator owns these bytes; the arena only retains and exposes it. + provider: Any + arena: Any + bound: Any + bundle: Any + plane_names: Tuple[str, ...] + local_plane_views: Tuple[torch.Tensor, ...] + tekit_alias_views: Tuple[torch.Tensor, ...] + home_experts: int + helper_slots: int + #: READY generation ``g`` currently published into this arena's helper + #: slots. 0 means "nothing published yet" -- the kernel ABI rejects it + #: (the gate requires g in [1, 2**63)), which is the intended fail-closed + #: state before the weight transport has run for this generation. + ready_generation_value: int = 0 + + @property + def slot_count(self) -> int: + return self.home_experts + self.helper_slots + + def terminal_flags_tensor(self) -> torch.Tensor: + """This rank's ``uint64[EP]`` terminal view -- the contract's + ``hot_expert_weight_ready_flags``. Sources publish into the multicast + alias of the SAME bytes, so a publish by any EP rank becomes visible + here without any host-side copy.""" + # Hierarchical arenas expose terminals through their base arena. + arena = getattr(self.arena, "base", self.arena) + return arena.terminals.local_view + + def ready_generation(self) -> int: + return int(self.ready_generation_value) + + def publish_ready_generation(self, generation: int) -> None: + """Record the generation the kernel must wait for. Called by the + weight transport AFTER it has issued the payload copies for ``g``; + the terminal cells themselves are written by the transport.""" + if type(generation) is not int or not (1 <= generation < (1 << 63)): + raise ValueError( + f"READY generation must be an exact int in [1, 2**63); got {generation!r}" + ) + self.ready_generation_value = generation + + def assert_identity(self) -> None: + """Run the optional allocator identity guard. + + The caller must also check framework parameter pointers: allocator checks + cannot detect replacement of a Parameter's storage. + """ + # The allocator guard is optional; framework pointer checks remain mandatory. + guard = getattr(self.bound, "assert_identity", None) + if guard is not None: + guard() + + +class _EpComm: + """Adapt the resolved torch EP ProcessGroup to SAMI's collectives. + + Using the default distributed group could include ranks outside this arena + and deadlock its collective allocation or binding. + """ + + def __init__(self, process_group: Any) -> None: + if process_group is None: + raise ValueError("MoE rebalance requires an EP ProcessGroup") + self.process_group = process_group + self.rank = int(torch.distributed.get_rank(group=process_group)) + self.world = int(torch.distributed.get_world_size(group=process_group)) + + def validate_geometry(self, *, expected_rank: int, expected_world: int) -> None: + expected = (int(expected_rank), int(expected_world)) + actual = (self.rank, self.world) + if actual != expected: + raise ValueError( + "EP ProcessGroup rank/size differs from the model mapping: " + f"process_group={actual}, mapping={expected}." + ) + + def barrier(self) -> None: + with _disable_current_modes(): + torch.distributed.barrier(group=self.process_group) + + def bcast(self, payload: object, root: int) -> object: + if type(root) is not int or not 0 <= root < self.world: + raise ValueError(f"broadcast root must be in [0, {self.world}); got {root!r}") + values = [payload] + source_rank = torch.distributed.get_global_rank(self.process_group, root) + with _disable_current_modes(): + torch.distributed.broadcast_object_list( + values, src=source_rank, group=self.process_group + ) + return values[0] + + def allgather(self, value: object) -> list: + values = [None] * self.world + with _disable_current_modes(): + torch.distributed.all_gather_object(values, value, group=self.process_group) + return values + + +class _RawPointer: + """Expose a CUDA address through __cuda_array_interface__. + + Tensor views do not own this allocation; the provider retains its owners. + """ + + def __init__(self, pointer: int, nbytes: int) -> None: + self.__cuda_array_interface__ = { + "data": (int(pointer), False), + "shape": (int(nbytes),), + "typestr": "|u1", + "strides": None, + "version": 3, + } + + +#: Planes whose whole [H + S] span is at most this many bytes (the fp32 scale +#: planes) live in a per-layer record of a small shared page instead of in +#: their own granularity-sized allocations. +_RECORD_PLANE_BYTES = 64 * 1024 +_RECORD_PLANE_ALIGNMENT = 256 + + +_HELPER_BANK_COUNT = 2 + + +class _ViewSpec: + """Address and layout of a plane view SAMI only binds, never dereferences. + + Peer and multicast views feed SAMI's bind-time pointer tables. A multicast + view's row 0 lies in a reserved but unmapped prefix, which a torch tensor + cannot wrap; only its helper rows are mapped and written by the copy. + """ + + __slots__ = ("_pointer", "shape", "_stride", "dtype", "_element_size", "device", "is_cuda") + + def __init__(self, pointer: int, like: torch.Tensor) -> None: + self._pointer = int(pointer) + self.shape = tuple(like.shape) + self._stride = tuple(like.stride()) + self.dtype = like.dtype + self._element_size = int(like.element_size()) + self.device = like.device + self.is_cuda = True + + def data_ptr(self) -> int: + return self._pointer + + def stride(self) -> Tuple[int, ...]: + return self._stride + + def element_size(self) -> int: + return self._element_size + + +class _Vmm: + """CUDA VMM calls on FABRIC-exportable memory with explicit ownership.""" + + def __init__(self, device: int, group_sizes: Tuple[int, ...]) -> None: + from cuda.bindings import driver as cuda + + from tensorrt_llm._torch.cute_dsl_kernels.megamoe_scheduler.sami._util import check_cuda + from tensorrt_llm._torch.cute_dsl_kernels.megamoe_scheduler.sami.fabric import ( + FABRIC, + _allocation_properties, + _read_write_access, + ) + + self.cuda = cuda + self.check = check_cuda + self.fabric = FABRIC + self.device = int(device) + self.properties = _allocation_properties(self.device) + self._access = [_read_write_access(self.device)] + # Every chunk is multicast-bindable at every level, so align to the + # coarsest MINIMUM granularity rather than the 512 MiB recommendation. + values = [ + BUFFER_ALIGNMENT, + int( + check_cuda( + cuda.cuMemGetAllocationGranularity( + self.properties, + cuda.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_MINIMUM, + ), + "cuMemGetAllocationGranularity(shared slots)", + ) + ), + ] + for members in sorted(set(group_sizes)): + values.append( + int( + check_cuda( + cuda.cuMulticastGetGranularity( + self.multicast_properties(members, BUFFER_ALIGNMENT), + cuda.CUmulticastGranularity_flags.CU_MULTICAST_GRANULARITY_MINIMUM, + ), + "cuMulticastGetGranularity(shared slots)", + ) + ) + ) + self.granularity = math.lcm(*values) + self.local_allocation_handles: List[Any] = [] + self.imported_allocation_handles: List[Any] = [] + self.created_multicast_handles: List[Any] = [] + self.imported_multicast_handles: List[Any] = [] + self.mappings: List[Tuple[int, int]] = [] + self.reservations: List[Tuple[int, int]] = [] + self.multicast_binds: List[Tuple[Any, int, int]] = [] + self.allocated_bytes = 0 + + def multicast_properties(self, members: int, nbytes: int) -> Any: + properties = self.cuda.CUmulticastObjectProp() + properties.numDevices = int(members) + properties.size = int(nbytes) + properties.handleTypes = int(self.fabric) + properties.flags = 0 + return properties + + def create(self, nbytes: int) -> Tuple[Any, bytes]: + handle = self.check( + self.cuda.cuMemCreate(nbytes, self.properties, 0), "cuMemCreate(shared slots)" + ) + self.local_allocation_handles.append(handle) + self.allocated_bytes += nbytes + exported = self.check( + self.cuda.cuMemExportToShareableHandle(handle, self.fabric, 0), + "cuMemExportToShareableHandle(shared slots)", + ) + shareable = bytes(exported.data) + if len(shareable) != 64: + raise RuntimeError(f"fabric handle has {len(shareable)} bytes; expected 64") + return handle, shareable + + def import_(self, shareable: bytes) -> Any: + if not isinstance(shareable, bytes) or len(shareable) != 64: + raise RuntimeError("peer fabric handle must contain exactly 64 bytes") + handle = self.check( + self.cuda.cuMemImportFromShareableHandle(bytearray(shareable), self.fabric), + "cuMemImportFromShareableHandle(shared slots)", + ) + self.imported_allocation_handles.append(handle) + return handle + + def register_created_multicast(self, handle: Any) -> None: + self.created_multicast_handles.append(handle) + + def import_multicast(self, shareable: bytes) -> Any: + if not isinstance(shareable, bytes) or len(shareable) != 64: + raise RuntimeError("multicast fabric handle must contain exactly 64 bytes") + handle = self.check( + self.cuda.cuMemImportFromShareableHandle(bytearray(shareable), self.fabric), + "cuMemImportFromShareableHandle(shared multicast)", + ) + self.imported_multicast_handles.append(handle) + return handle + + def reserve(self, nbytes: int) -> int: + address = int( + self.check( + self.cuda.cuMemAddressReserve(nbytes, self.granularity, 0, 0), + "cuMemAddressReserve(shared slots)", + ) + ) + self.reservations.append((address, int(nbytes))) + return address + + def map(self, address: int, nbytes: int, handle: Any) -> None: + self.check(self.cuda.cuMemMap(address, nbytes, 0, handle, 0), "cuMemMap(shared slots)") + self.mappings.append((int(address), int(nbytes))) + self.check( + self.cuda.cuMemSetAccess(address, nbytes, self._access, 1), + "cuMemSetAccess(shared slots)", + ) + + def record_multicast_bind(self, handle: Any, offset: int, nbytes: int) -> None: + self.multicast_binds.append((handle, int(offset), int(nbytes))) + + def unmap_all(self) -> None: + while self.mappings: + address, nbytes = self.mappings[-1] + self.check(self.cuda.cuMemUnmap(address, nbytes), "cuMemUnmap(shared slots)") + self.mappings.pop() + + def unbind_all(self) -> None: + while self.multicast_binds: + handle, offset, nbytes = self.multicast_binds[-1] + self.check( + self.cuda.cuMulticastUnbind(handle, self.device, offset, nbytes), + "cuMulticastUnbind(shared slots)", + ) + self.multicast_binds.pop() + + def free_reservations(self) -> None: + while self.reservations: + address, nbytes = self.reservations[-1] + self.check( + self.cuda.cuMemAddressFree(address, nbytes), + "cuMemAddressFree(shared slots)", + ) + self.reservations.pop() + + def _release_handles(self, handles: List[Any], label: str) -> None: + while handles: + self.check(self.cuda.cuMemRelease(handles[-1]), label) + handles.pop() + + def release_imported_handles(self) -> None: + self._release_handles( + self.imported_multicast_handles, + "cuMemRelease(imported shared multicast)", + ) + self._release_handles( + self.imported_allocation_handles, + "cuMemRelease(imported shared allocation)", + ) + + def release_created_handles(self) -> None: + self._release_handles( + self.created_multicast_handles, + "cuMemRelease(created shared multicast)", + ) + self._release_handles( + self.local_allocation_handles, + "cuMemRelease(local shared allocation)", + ) + + def zero(self, address: int, nbytes: int) -> None: + self.check( + self.cuda.cuMemsetD8(self.cuda.CUdeviceptr(address), 0, nbytes), + "cuMemsetD8(shared slots)", + ) + + +class _MulticastSpec: + """One aligned group's multicast pair: member binds and local mappings.""" + + def __init__( + self, + members: Tuple[int, ...], + nbytes: int, + binds: List[Tuple[int, Any, int]], + mappings: List[int], + ) -> None: + self.members = members + self.nbytes = nbytes + #: (multicast offset, member allocation handle, bytes) bound on every member. + self.binds = binds + #: Reserved addresses where this rank maps its selected object. + self.mappings = mappings + self.handles: List[Any] = [] + + +def _create_multicast(vmm: _Vmm, comm: Any, rank: int, specs: List[_MulticastSpec]) -> None: + """Collectively create, bind and map dual-creator multicast objects. + + Follows MulticastTeamPair: the first two members each create one object, + every member binds its chunks to both, and the first member maps the second + member's object while every other member maps the first member's. Ranks + outside a group join only the collectives; specs are rank-ordered alike. + """ + cuda, check = vmm.cuda, vmm.check + created: Dict[Tuple[int, int], Any] = {} + exported: Dict[Tuple[int, int], bytes] = {} + for index, spec in enumerate(specs): + for alias, creator in enumerate(spec.members[:2]): + if rank == creator: + handle = check( + cuda.cuMulticastCreate( + vmm.multicast_properties(len(spec.members), spec.nbytes) + ), + "cuMulticastCreate(shared slots)", + ) + vmm.register_created_multicast(handle) + created[(index, alias)] = handle + exported[(index, alias)] = bytes( + check( + cuda.cuMemExportToShareableHandle(handle, vmm.fabric, 0), + "cuMemExportToShareableHandle(shared multicast)", + ).data + ) + gathered = comm.allgather(exported) + joined = [(index, spec) for index, spec in enumerate(specs) if rank in spec.members] + for index, spec in joined: + for alias, creator in enumerate(spec.members[:2]): + handle = created.get((index, alias)) + if handle is None: + handle = vmm.import_multicast(gathered[creator][(index, alias)]) + spec.handles.append(handle) + check( + cuda.cuMulticastAddDevice(handle, vmm.device), "cuMulticastAddDevice(shared slots)" + ) + comm.barrier() + for _, spec in joined: + for handle in spec.handles: + for offset, memory, nbytes in spec.binds: + check( + cuda.cuMulticastBindMem_v2(handle, vmm.device, offset, memory, 0, nbytes, 0), + "cuMulticastBindMem_v2(shared slots)", + ) + vmm.record_multicast_bind(handle, offset, nbytes) + comm.barrier() + for _, spec in joined: + selected = spec.handles[1 if len(spec.handles) > 1 and rank == spec.members[0] else 0] + for address in spec.mappings: + vmm.map(address, spec.nbytes, selected) + comm.barrier() + + +@dataclass +class _HelperSet: + """One physical helper-slot set: a chunk per large plane, one multicast + object per hierarchy level, and that object's view base per plane.""" + + chunks: Dict[str, Any] + level_bases: Tuple[Dict[str, int], ...] + level_specs: Tuple[_MulticastSpec, ...] + zeroed: bool = False + + +@dataclass +class _RecordPage: + """Per-layer records of READY terminals and small planes, mapped on every + rank and bound to one multicast object per hierarchy level.""" + + uc_ptrs: Tuple[int, ...] + mc_ptrs: Tuple[int, ...] + level_specs: Tuple[_MulticastSpec, ...] + + +class _SharedLayer: + """Anti-GC owner of one layer's home chunks and raw pointer holders.""" + + def __init__(self, pool: "_SharedSlotPool", helper_set: int) -> None: + self.pool = pool + self.helper_set = helper_set + self.home_handles: List[Any] = [] + self.holders: List[_RawPointer] = [] + + +class _SharedSlotPool: + """Helper-slot sets shared by every MoE layer of one geometry. + + Layer ``i`` maps helper set ``i % sets`` directly after its home rows, so + each plane is still one contiguous ``[H + S]`` kernel view. Only the home + rows are per layer. No extra synchronization guards reuse: layer ``i``'s + copy follows HALO-Q's all-rank exchange, which every rank enters only after + MAIN finished the previous layers' MegaMoE. + + READY terminals stay per layer because one forward shares its generation + across layers. They and the fp32 planes live in per-layer records. + """ + + def __init__( + self, + *, + world: int, + rank: int, + home_count: int, + helper_count: int, + bundle: Any, + device: int, + comm: Any, + owner_registry: List["_SharedSlotPool"], + ) -> None: + from tensorrt_llm._torch.cute_dsl_kernels.megamoe_scheduler.sami import ( + hierarchy_group_sizes, + ) + + torch.cuda.set_device(device) + self.world = world + self.rank = rank + self.device = device + self.comm = comm + self._closed = False + self._closing = False + self._producers: Dict[int, Any] = {} + self.home_count = home_count + self.helper_count = helper_count + self.slots = home_count + helper_count + self.bundle = bundle + self.group_sizes = hierarchy_group_sizes(world) + self.groups = tuple( + tuple(range(rank // size * size, rank // size * size + size)) + for size in self.group_sizes + ) + self.vmm = _Vmm(device, self.group_sizes) + granularity = self.vmm.granularity + + self.large = tuple(p for p in bundle.planes if p.nbytes * self.slots > _RECORD_PLANE_BYTES) + self.home_bytes = { + p.name: _align_up(home_count * p.nbytes, granularity) for p in self.large + } + self.helper_bytes = { + p.name: _align_up(helper_count * p.nbytes, granularity) for p in self.large + } + # One multicast object per set and level holds every large plane's chunk. + self.mc_offsets: Dict[str, int] = {} + cursor = 0 + for plane in self.large: + self.mc_offsets[plane.name] = cursor + cursor += self.helper_bytes[plane.name] + self.set_bytes = cursor + + # A record keeps its READY terminals on their own page. + cursor = _align_up(world * 8, TERMINAL_ALIGNMENT) + self.record_offsets: Dict[str, int] = {} + for plane in bundle.planes: + if plane.name not in self.home_bytes: + cursor = _align_up(cursor, _RECORD_PLANE_ALIGNMENT) + self.record_offsets[plane.name] = cursor + cursor += self.slots * plane.nbytes + self.record_stride = _align_up(cursor, TERMINAL_ALIGNMENT) + self.page_bytes = _align_up(self.record_stride, granularity) + self.records_per_page = self.page_bytes // self.record_stride + self.pages: List[_RecordPage] = [] + self.layer_count = 0 + self.sets: Tuple[_HelperSet, ...] = () + self._construction_complete = False + self._owner_registry = owner_registry + self._owner_registry.append(self) + self.sets = self._create_sets(comm) + self._construction_complete = True + + def register_producer(self, producer: Any) -> None: + if self._closed: + raise RuntimeError("cannot register a producer on a closed shared-slot pool") + self._producers.setdefault(id(producer), producer) + + def _cleanup_manifest(self) -> tuple: + return ( + 1, + self._construction_complete, + self.world, + self.home_count, + self.helper_count, + self.bundle.hidden, + self.bundle.intermediate, + self.group_sizes, + self.layer_count, + len(self.pages), + len(self.sets), + tuple((plane.name, int(plane.nbytes)) for plane in self.bundle.planes), + tuple( + ( + type(producer).__qualname__, + getattr(producer, "layer_idx", getattr(producer, "_layer_idx", None)), + ) + for producer in self._producers.values() + ), + ) + + def _collective_phase(self, label: str, operation: Any) -> None: + error = None + try: + operation() + except Exception as exc: + error = f"{type(exc).__name__}: {exc}" + errors = self.comm.allgather(error) + if len(errors) != self.world: + raise RuntimeError( + f"shared-slot cleanup {label} returned {len(errors)} ranks; expected {self.world}" + ) + failures = [f"rank {rank}: {item}" for rank, item in enumerate(errors) if item] + if failures: + raise RuntimeError(f"shared-slot cleanup {label} failed: {'; '.join(failures)}") + + def _close_producers(self) -> None: + for producer in self._producers.values(): + close = getattr(producer, "close", None) + if close is None: + raise TypeError(f"registered producer {type(producer).__qualname__} has no close") + if close() is False: + raise RuntimeError( + f"registered producer {type(producer).__qualname__} did not close" + ) + + def _on_device(self, operation: Any) -> None: + with torch.cuda.device(self.device): + operation() + + def close_collectively(self, local_safe: bool = True) -> bool: + if self._closed: + return True + if self._closing: + raise RuntimeError("shared-slot pool cleanup is already in progress") + preflight = self.comm.allgather( + { + "local_safe": bool(local_safe), + "manifest": self._cleanup_manifest(), + }, + ) + if len(preflight) != self.world: + raise RuntimeError( + f"shared-slot cleanup preflight returned {len(preflight)} ranks; " + f"expected {self.world}" + ) + if not all(item["local_safe"] for item in preflight): + return False + manifests = [item["manifest"] for item in preflight] + expected = manifests[0] + if any(manifest != expected for manifest in manifests[1:]): + raise RuntimeError("shared-slot cleanup manifests differ across the EP communicator") + self._closing = True + try: + self._collective_phase( + "device synchronize", + lambda: self._on_device(lambda: torch.cuda.synchronize(self.device)), + ) + self._collective_phase("producers", lambda: self._on_device(self._close_producers)) + self._producers.clear() + self._collective_phase("unmap", lambda: self._on_device(self.vmm.unmap_all)) + self._collective_phase("multicast unbind", lambda: self._on_device(self.vmm.unbind_all)) + self._collective_phase( + "address free", lambda: self._on_device(self.vmm.free_reservations) + ) + self._collective_phase( + "release imported handles", + lambda: self._on_device(self.vmm.release_imported_handles), + ) + self._collective_phase( + "release created handles", + lambda: self._on_device(self.vmm.release_created_handles), + ) + self._closed = True + if self in self._owner_registry: + self._owner_registry.remove(self) + self.sets = () + self.pages.clear() + return True + finally: + self._closing = False + + def _create_sets(self, comm: Any) -> Tuple[_HelperSet, ...]: + if not self.large: + return tuple( + _HelperSet({}, tuple({} for _ in self.groups), ()) + for _ in range(_HELPER_BANK_COUNT) + ) + sets, specs = [], [] + for _ in range(_HELPER_BANK_COUNT): + chunks = {p.name: self.vmm.create(self.helper_bytes[p.name])[0] for p in self.large} + level_bases, level_specs = [], [] + for members in self.groups: + bases: Dict[str, int] = {} + mappings: List[int] = [] + for plane in self.large: + # Row 0 of the view sits H rows before the plane's chunk. + prefix = self.home_bytes[plane.name] + address = self.vmm.reserve(prefix + self.set_bytes) + mappings.append(address + prefix) + bases[plane.name] = ( + address + + prefix + + self.mc_offsets[plane.name] + - self.home_count * plane.nbytes + ) + spec = _MulticastSpec( + members, + self.set_bytes, + [ + (self.mc_offsets[p.name], chunks[p.name], self.helper_bytes[p.name]) + for p in self.large + ], + mappings, + ) + specs.append(spec) + level_specs.append(spec) + level_bases.append(bases) + sets.append(_HelperSet(chunks, tuple(level_bases), tuple(level_specs))) + _create_multicast(self.vmm, comm, self.rank, specs) + return tuple(sets) + + def _create_page(self, comm: Any) -> _RecordPage: + handle, shareable = self.vmm.create(self.page_bytes) + local = self.vmm.reserve(self.page_bytes) + self.vmm.map(local, self.page_bytes, handle) + self.vmm.zero(local, self.page_bytes) + torch.cuda.synchronize(self.device) + shareables = comm.allgather(shareable) + if len(shareables) != self.world: + raise RuntimeError("shared-slot record exchange returned an incomplete EP") + uc_ptrs = [] + for peer in range(self.world): + if peer == self.rank: + uc_ptrs.append(local) + continue + address = self.vmm.reserve(self.page_bytes) + self.vmm.map(address, self.page_bytes, self.vmm.import_(shareables[peer])) + uc_ptrs.append(address) + mc_ptrs = [self.vmm.reserve(self.page_bytes) for _ in self.groups] + specs = [ + _MulticastSpec(members, self.page_bytes, [(0, handle, self.page_bytes)], [address]) + for members, address in zip(self.groups, mc_ptrs) + ] + _create_multicast(self.vmm, comm, self.rank, specs) + return _RecordPage(tuple(uc_ptrs), tuple(mc_ptrs), tuple(specs)) + + def _flat(self, owner: _SharedLayer, pointer: int, nbytes: int) -> torch.Tensor: + holder = _RawPointer(pointer, nbytes) + owner.holders.append(holder) + return torch.as_tensor(holder, device=f"cuda:{self.device}") + + def build_layer(self, comm: Any) -> Tuple[Any, Dict[str, torch.Tensor], _SharedLayer, int]: + """Collectively map one layer; returns its arena, framework views, + owner and own allocation bytes.""" + from tensorrt_llm._torch.cute_dsl_kernels.megamoe_scheduler.sami import ( + HierarchicalLiveWeightArena, + LivePlaneView, + LiveTerminalView, + LiveWeightArena, + ) + + world, rank, slots = self.world, self.rank, self.slots + index = self.layer_count + if index % self.records_per_page == 0: + self.pages.append(self._create_page(comm)) + page = self.pages[-1] + record = (index % self.records_per_page) * self.record_stride + helper_set = self.sets[index % len(self.sets)] + layer = _SharedLayer(self, index % len(self.sets)) + self.layer_count += 1 + + homes = {p.name: self.vmm.create(self.home_bytes[p.name]) for p in self.large} + layer.home_handles = [homes[p.name][0] for p in self.large] + peer_handles = comm.allgather(tuple(homes[p.name][1] for p in self.large)) + if len(peer_handles) != world: + raise RuntimeError("shared-slot home exchange returned an incomplete EP") + + planes, group_views, group_owners = [], [], [] + tekit_views: Dict[str, torch.Tensor] = {} + for plane in self.bundle.planes: + span = slots * plane.nbytes + if plane.name in self.home_bytes: + large_index = [p.name for p in self.large].index(plane.name) + home_bytes = self.home_bytes[plane.name] + helper_bytes = self.helper_bytes[plane.name] + # Home rows end where the helper chunk begins. + pad = home_bytes - self.home_count * plane.nbytes + address = self.vmm.reserve(home_bytes + helper_bytes) + self.vmm.map(address, home_bytes, homes[plane.name][0]) + self.vmm.map(address + home_bytes, helper_bytes, helper_set.chunks[plane.name]) + self.vmm.zero(address, home_bytes) + if not helper_set.zeroed: + self.vmm.zero(address + home_bytes, helper_bytes) + uc_ptrs = [] + for peer in range(world): + if peer == rank: + uc_ptrs.append(address + pad) + continue + # Peers expose only home rows; the reservation keeps the + # view's helper tail clear of every other mapping. + peer_address = self.vmm.reserve(home_bytes + helper_bytes) + self.vmm.map( + peer_address, + home_bytes, + self.vmm.import_(peer_handles[peer][large_index]), + ) + uc_ptrs.append(peer_address + pad) + level_ptrs = tuple(bases[plane.name] for bases in helper_set.level_bases) + owners = helper_set.level_specs + else: + offset = record + self.record_offsets[plane.name] + uc_ptrs = [pointer + offset for pointer in page.uc_ptrs] + level_ptrs = tuple(pointer + offset for pointer in page.mc_ptrs) + owners = page.level_specs + flat = self._flat(layer, uc_ptrs[rank], span) + kernel_view = _plane_view(flat, plane, slots, kernel=True) + uc_views = tuple( + kernel_view if peer == rank else _ViewSpec(pointer, kernel_view) + for peer, pointer in enumerate(uc_ptrs) + ) + level_views = tuple(_ViewSpec(pointer, kernel_view) for pointer in level_ptrs) + tekit_views[plane.name] = _plane_view(flat, plane, slots, kernel=False) + planes.append( + LivePlaneView( + name=plane.name, + uc_views=uc_views, + mc_view=level_views[0], + aliases_same_backing=True, + backing_owner=layer, + ) + ) + group_views.append(level_views) + group_owners.append(tuple(owners)) + helper_set.zeroed = True + + terminals = LiveTerminalView( + local_view=self._flat(layer, page.uc_ptrs[rank] + record, world * 8).view(torch.uint64), + mc_view=self._flat(layer, page.mc_ptrs[0] + record, world * 8).view(torch.uint64), + aliases_same_backing=True, + backing_owner=layer, + ) + arena = HierarchicalLiveWeightArena( + base=LiveWeightArena(planes=tuple(planes), terminals=terminals), + group_sizes=self.group_sizes, + group_mc_views=tuple(group_views), + group_mc_owners=tuple(group_owners), + aliases_same_backing=True, + ) + return arena, tekit_views, layer, sum(self.home_bytes.values()) + + +class SharedSlotArenaProvider: + """Provide one layer's arena over the two shared helper-slot banks. + + The broadcaster requests the arena again after Parameters are bound, so a + provider memoizes its first build and requires identical geometry. + """ + + def __init__(self, owner: Any) -> None: + # The model-level per-layer dynamic manager owns both the geometry cache + # and creation order. No Mapping attribute or process-global registry + # participates in allocation lifetime. + self._owner = owner + pools = getattr(owner, "_rebalance_shared_slot_pools", None) + order = getattr(owner, "_rebalance_shared_slot_pool_order", None) + if not isinstance(pools, dict) or not isinstance(order, list): + raise TypeError("SharedSlotArenaProvider owner must expose model-local pool state") + self._pools = pools + self._pool_order = order + self.buffer_bytes = 0 + self.group_sizes: Tuple[int, ...] = () + self.pool: Optional[_SharedSlotPool] = None + self.layer: Optional[_SharedLayer] = None + self._built: Optional[Any] = None + self._signature: Optional[tuple] = None + self._tekit_views: Dict[str, torch.Tensor] = {} + + def build_hierarchical_live_arena( + self, + *, + world: int, + rank: int, + home_count: int, + helper_count: int, + bundle: Any, + device: int, + comm: Any, + ) -> Any: + signature = ( + int(world), + int(rank), + int(home_count), + int(helper_count), + int(device), + bundle, + ) + if self._built is not None: + if self._signature != signature: + raise RuntimeError( + "shared-slot arena provider is per-layer and single-geometry; " + f"rebuilt with {signature} after {self._signature}" + ) + if self.pool is not None and not self.pool._closed: + return self._built + self._built = None + self._tekit_views = {} + self.layer = None + self.pool = None + self.buffer_bytes = 0 + self.group_sizes = () + self._signature = signature + key = signature + pool = self._pools.get(key) + if pool is not None and pool._closed: + self._pools.pop(key) + pool = None + if pool is None: + pool = _SharedSlotPool( + world=int(world), + rank=int(rank), + home_count=int(home_count), + helper_count=int(helper_count), + bundle=bundle, + device=int(device), + comm=comm, + owner_registry=self._pool_order, + ) + self._pools[key] = pool + self._built, self._tekit_views, self.layer, self.buffer_bytes = pool.build_layer(comm) + self.pool = pool + self.group_sizes = pool.group_sizes + return self._built + + def register_producer(self, producer: Any) -> None: + if self.pool is None: + raise RuntimeError("build the shared-slot arena before registering its producer") + self.pool.register_producer(producer) + + def tekit_alias(self, name: str) -> torch.Tensor: + """Return this rank's framework storage view for a weight plane.""" + return self._tekit_views[name] + + +def allocate_rebalance_arena( + *, + home_experts: int, + helper_slots: int, + bundle: Any, + ep_comm: Any, + mapping: Any, + provider: SharedSlotArenaProvider, + device: int, + layer_idx: Optional[int], +) -> RebalanceLiveArena: + """Allocate and bind one layer's arena collectively across its EP group. + + Every rank must call this with identical geometry and collective order. + """ + world = int(mapping.moe_ep_size) + rank = int(mapping.moe_ep_rank) + if world < 2: + raise RuntimeError(f"MoE rebalance needs at least two EP ranks; got moe_ep_size={world}.") + + if bundle is None or not getattr(bundle, "planes", None): + raise TypeError("rebalance arena requires a quant-method live-weight plane spec") + if not isinstance(ep_comm, _EpComm): + raise TypeError("rebalance arena requires its model-owned EP communicator") + ep_comm.validate_geometry(expected_rank=rank, expected_world=world) + if not isinstance(provider, SharedSlotArenaProvider): + raise TypeError("rebalance arena provider must be owned by the model load balancer") + arena = provider.build_hierarchical_live_arena( + world=world, + rank=rank, + home_count=int(home_experts), + helper_count=int(helper_slots), + bundle=bundle, + device=int(device), + comm=ep_comm, + ) + # Initialize READY terminals explicitly before peers can observe them. + # This startup-only reset does not depend on allocator initialization details. + arena.base.terminals.local_view.zero_() + torch.cuda.synchronize() + ep_comm.barrier() + bound = arena.bind( + world=world, + rank=rank, + home_count=int(home_experts), + helper_count=int(helper_slots), + bundle=bundle, + device=int(device), + ) + plane_names = tuple(plane.name for plane in bundle.planes) + # Retain the hierarchical arena and its bound base views. + live = RebalanceLiveArena( + provider=provider, + arena=arena, + bound=bound, + bundle=bundle, + plane_names=plane_names, + # The bound hierarchy exposes local plane views through its base arena. + local_plane_views=tuple(bound.base.local_plane_views), + tekit_alias_views=tuple(provider.tekit_alias(name) for name in plane_names), + home_experts=int(home_experts), + helper_slots=int(helper_slots), + ) + pool = provider.pool + layer = provider.layer + assert pool is not None and layer is not None + logger.debug( + f"[MegaMoECuteDsl] layer={layer_idx} MoE rebalance SHARED-SLOT live arena: " + f"{provider.buffer_bytes} B of home rows for M={live.slot_count} slots " + f"(H={home_experts} + S={helper_slots}) on helper set " + f"{layer.helper_set}/{len(pool.sets)}; pool sets " + f"{len(pool.sets)} x {pool.set_bytes} B, record pages {len(pool.pages)} x " + f"{pool.page_bytes} B, granularity {pool.vmm.granularity} B, VMM total " + f"{pool.vmm.allocated_bytes} B after {pool.layer_count} layers; ep_size={world} " + f"ep_rank={rank} device={device}, group_sizes={provider.group_sizes}." + ) + return live diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py new file mode 100644 index 000000000000..2ffa6745cb55 --- /dev/null +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py @@ -0,0 +1,503 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Single-main-thread HALO-Q and in-switch TMA weight-copy integration. + +One highest-priority stream is shared across layers. plan_schedule() enqueues +HALO-Q, records route completion, and enqueues TMA copy before returning to the +shared-expert hook. plan_finish() can lend routes before the route event; +the first route consumer calls wait_for_routes(). MegaMoE's +system-acquire READY gate protects helper reads while TMA finishes concurrently. + +For each rank and generation g: + MAIN: MegaMoE(g) -> consumer_done(g) + COPY: wait consumer_done(g) -> HALO-Q(g+1) EP exchange -> TMA(g+1) +All ranks enter the bounded EP exchange after their own previous consumer, so no +next-generation overwrite can precede any previous-generation consumer. No host +collective or auxiliary copy submitter is needed in the serving path. +""" + +from __future__ import annotations + +import threading +from typing import Any, Optional, Tuple + +import torch + +from tensorrt_llm.logger import logger + +__all__ = ["RebalanceSlotSchedulerGroup", "build_rebalance_slot_scheduler_group"] + +# Shared-expert persistent grids use the same capacity hint as the copy grid. +TMA_COPY_SM_COUNT = 8 + + +class _LiveBankLeaseProvider: + """Authorize ordered reuse, not a claim of completed GPU work on the host. + + release_generation_after() enqueues the local consumer wait on COPY. + prove() is legal only after the next HALO-Q launch on that same stream. + Its all-rank count exchange is bounded and precedes plan publication; TMA + waits for that publication and executes after HALO-Q. Therefore a delayed + peer consumer blocks every rank's overwrite without a host-side barrier. + """ + + collective_safe = True + producer_reuse_guarded = True + live_bank_count = 1 + + def __init__(self, group: RebalanceSlotSchedulerGroup) -> None: + self.bound_copy_module = group.broadcaster + self._group = group + self.proved_generation = 0 + + def prove(self, generation: int) -> None: + group = self._group + if ( + generation != group._finished_generation + or group._pending_release is None + or group._pending_release[0] != generation + or group._scheduled_generation != generation + 1 + or group.scheduler_stream is not group.copy_stream + ): + raise RuntimeError("Live-bank reuse lacks the consumer/HALO-Q ordering chain") + self.bound_copy_module.mark_collective_reuse_safe(self, generation) + self.proved_generation = generation + + +class RebalanceSlotSchedulerGroup: + """One layer's scheduler, copy endpoint, and paired consumer generations.""" + + def __init__( + self, + *, + arena: Any, + mapping: Any, + ep_comm: Any, + device: int, + home_experts: int, + helper_slots: int, + topk: int, + max_tokens_per_rank: int, + copy_stream: torch.cuda.Stream, + auxiliary_sms: int, + layer_idx: Optional[int], + ) -> None: + from cuda.bindings import driver + + from tensorrt_llm._torch.cute_dsl_kernels.megamoe_scheduler.cuda_scheduler import ( + CudaPhysicalSlotScheduler, + CudaSchedulerConfig, + ) + from tensorrt_llm._torch.cute_dsl_kernels.megamoe_scheduler.cuda_scheduler.multirank_dist import ( + FabricSymmetricBuffer, + ) + from tensorrt_llm._torch.cute_dsl_kernels.megamoe_scheduler.cuda_scheduler.runtime import ( + symmetric_buffer_ints, + ) + from tensorrt_llm._torch.cute_dsl_kernels.megamoe_scheduler.sami import ( + HierarchicalSamiWeightBroadcast, + ) + + ep_size = int(mapping.moe_ep_size) + ep_rank = int(mapping.moe_ep_rank) + self.layer_idx = layer_idx + self.device = int(device) + self.ep_size = ep_size + self.ep_rank = ep_rank + self.home_experts = int(home_experts) + self.helper_slots = int(helper_slots) + self.topk = int(topk) + self.max_tokens_per_rank = int(max_tokens_per_rank) + self.auxiliary_sms = int(auxiliary_sms) + if self.auxiliary_sms < TMA_COPY_SM_COUNT: + raise ValueError( + f"auxiliary_sms must be at least {TMA_COPY_SM_COUNT}; got {self.auxiliary_sms}" + ) + if copy_stream.device.index != self.device: + raise ValueError("copy stream device does not match scheduler device") + self.arena = arena + self.plan_calls = 0 + self._plan_part: Optional[Tuple[Any, int]] = None + self._generation = 0 + self._pending_release: Optional[Tuple[int, Any]] = None + self._finished_generation = 0 + self._scheduled_generation = 0 + self._route_wait_generation = 0 + self._owner_thread_id: Optional[int] = None + self._execution_stream_handle: Optional[int] = None + self._closed = False + self.scheduler = None + self.symmetric = None + self.broadcaster = None + self.lease = None + + # Keep a partially initialized producer reachable by the arena's + # collective shutdown path. Constructor failures must not attempt a + # rank-local teardown while peers may still be initializing. + ep_comm.validate_geometry(expected_rank=ep_rank, expected_world=ep_size) + self._ep_comm = ep_comm + pool = getattr(arena.provider, "pool", None) + if pool is None or pool.comm is not ep_comm: + raise RuntimeError("rebalance arena and scheduler must use the same EP ProcessGroup") + arena.provider.register_producer(self) + + logical_expert_count = self.home_experts * ep_size + + # Use the EP rank, not the CUDA ordinal. HALO-Q shares the configured + # copy-stream launch budget with TMA. + cfg = CudaSchedulerConfig( + ep_size=ep_size, + logical_expert_count=logical_expert_count, + extra_slots_per_rank=self.helper_slots, + max_tokens_per_rank=self.max_tokens_per_rank, + topk=self.topk, + local_rank=ep_rank, + ctas=self.auxiliary_sms, + ) + cfg.validate() + self.cfg = cfg + self.scheduler = CudaPhysicalSlotScheduler(cfg, device=f"cuda:{self.device}") + + # All layers share the model manager's highest-priority copy stream. + self.copy_stream = copy_stream + self.scheduler_stream = self.copy_stream + self._stream_handle = int(self.copy_stream.cuda_stream) + self._driver = driver + self._input_ready = torch.cuda.Event() + self._plan_ready = torch.cuda.Event() + self._consumer_done = torch.cuda.Event() + # Materialize each lazily created event once. The steady path uses its + # cached driver handle, without per-plan stream guards or allocation. + with torch.cuda.device(self.device): + for event in (self._input_ready, self._plan_ready, self._consumer_done): + event.record(self.copy_stream) + self._input_ready_handle = int(self._input_ready.cuda_event) + self._plan_ready_handle = int(self._plan_ready.cuda_event) + self._consumer_done_handle = int(self._consumer_done.cuda_event) + + # Every peer address must be connected before either warmup or submit; + # HALO-Q performs an all-rank exchange on the COPY stream. + self.symmetric = FabricSymmetricBuffer( + 4 * symmetric_buffer_ints(ep_size, logical_expert_count), self.device + ) + handles = list(ep_comm.allgather(self.symmetric.shareable)) + if len(handles) != ep_size or len(set(handles)) != ep_size: + raise RuntimeError( + f"MoE rebalance requires {ep_size} unique EP fabric handles; " + f"received {len(handles)} handles with {len(set(handles))} unique values." + ) + self.scheduler.connect_peer_bases( + [ + self.symmetric.ptr if r == ep_rank else self.symmetric.import_peer(handles[r]) + for r in range(ep_size) + ] + ) + + # Warm all ranks before binding the plan channel: a bound launch would + # leave a pending publication that must be consumed by the broadcaster. + # The scheduler input is already initialized and peer allocations synchronized. + self.scheduler.launch(self.scheduler_stream) + torch.cuda.synchronize(self.device) + ep_comm.barrier() + + # The provider must return the existing arena with identical bundle geometry. + # A second allocation would separate the loader aliases from the copy target. + bundle = arena.bundle + self.broadcaster = HierarchicalSamiWeightBroadcast( + comm=ep_comm, + device=self.device, + bundle=bundle, + helper_count=self.helper_slots, + global_expert_count=logical_expert_count, + arena_provider=arena.provider, + tma_sm_count=TMA_COPY_SM_COUNT, + tma_warps=7, + tma_route="plan", + ) + self.broadcaster.bind_scheduler(self.scheduler, self.scheduler_stream) + + # A prior consumer event precedes the next HALO-Q EP rendezvous. + self.lease = _LiveBankLeaseProvider(self) + self.broadcaster.bind_generation_reuse_authority(self.lease) + + logger.debug("[MegaMoECuteDsl] layer=%s rebalance producer initialized", layer_idx) + + def close(self) -> None: + """Release this layer after the owning pool has drained device work.""" + if self._closed: + return + with torch.cuda.device(self.device): + if self.broadcaster is not None: + self.broadcaster.close() + if self.symmetric is not None: + self.symmetric.close() + self._closed = True + self._owner_thread_id = None + self._execution_stream_handle = None + self._plan_part = None + self._pending_release = None + self.lease = None + self.broadcaster = None + self.scheduler = None + self.symmetric = None + self.arena = None + + @property + def has_submission_owner(self) -> bool: + """Whether a CPU thread currently owns MAIN submissions.""" + return self._owner_thread_id is not None + + def _check_owner(self) -> None: + if self._closed: + raise RuntimeError("Rebalance scheduler group is closed") + thread_id = threading.get_ident() + if torch.cuda.current_device() != self.device: + raise RuntimeError("Rebalance submission requires its bound CUDA device") + stream = torch.cuda.current_stream(self.device) + handle = int(stream.cuda_stream) + if self._owner_thread_id is None: + if ( + self._execution_stream_handle is not None + and handle != self._execution_stream_handle + ): + raise RuntimeError("Rebalance MAIN stream cannot change during owner handoff") + if handle == self._stream_handle or stream.priority <= self.copy_stream.priority: + raise RuntimeError("MAIN must be distinct from the higher-priority COPY stream") + self._owner_thread_id = thread_id + self._execution_stream_handle = handle + elif thread_id != self._owner_thread_id or handle != self._execution_stream_handle: + raise RuntimeError("Rebalance must be submitted by one MAIN thread and stream") + + def release_warmup_owner_after_sync(self, execution_stream_handle: int) -> None: + """Release the startup CPU owner after the caller drains this device. + + Keep the fixed MAIN stream, generation counters, and pending lease + release. The first serving submission binds its CPU owner normally. + """ + if self._closed: + raise RuntimeError("Rebalance scheduler group is closed") + if self._owner_thread_id != threading.get_ident(): + raise RuntimeError("Only the warmup submitter may release its owner") + if self._execution_stream_handle != int(execution_stream_handle): + raise RuntimeError("Warmup owner handoff requires the bound MAIN stream") + if ( + self._plan_part is not None + or self._generation != self._finished_generation + or self._generation != self.plan_calls + ): + raise RuntimeError("Warmup owner handoff requires all plans to be finished") + self._owner_thread_id = None + + def release_quiesced_submission_owner(self, execution_stream_handle: int) -> None: + """Release a stopped executor's MAIN owner after a device drain. + + Unlike the warmup handoff, this is called by the shutdown thread after + the serving worker has exited. The caller must synchronize the device + first, so neither the old CPU thread nor its MAIN stream can still be + submitting work. + """ + if self._closed: + raise RuntimeError("Rebalance scheduler group is closed") + handle = int(execution_stream_handle) + if ( + self._plan_part is not None + or self._generation != self._finished_generation + or self._generation != self.plan_calls + ): + raise RuntimeError("Quiesced owner release requires all plans to be finished") + if self._execution_stream_handle is not None and self._execution_stream_handle != handle: + raise RuntimeError("Quiesced owner release requires the bound MAIN stream") + self._owner_thread_id = None + self._execution_stream_handle = None + + def _record_event(self, event_handle: int, stream_handle: int) -> None: + (error,) = self._driver.cuEventRecord(event_handle, stream_handle) + if error != self._driver.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"cuEventRecord(rebalance) failed: {error!r}") + + def _wait_event(self, stream_handle: int, event_handle: int) -> None: + (error,) = self._driver.cuStreamWaitEvent(stream_handle, event_handle, 0) + if error != self._driver.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"cuStreamWaitEvent(rebalance) failed: {error!r}") + + def plan_schedule(self, logical_expert_ids: torch.Tensor) -> Tuple[Any, int]: + """Submit contiguous CUDA int32 routes [tokens, topk] without staging. + + The route producer must run on MAIN and must not overwrite this tensor + until HALO-Q finishes. record_stream protects allocator reuse while COPY + consumes the input. Empty ranks still submit and join the EP exchange. + """ + self._check_owner() + if self._plan_part is not None or self._generation != self._finished_generation: + raise RuntimeError("Previous rebalance plan has no paired finish") + if ( + not isinstance(logical_expert_ids, torch.Tensor) + or logical_expert_ids.ndim != 2 + or logical_expert_ids.shape[1] != self.topk + or logical_expert_ids.dtype != torch.int32 + or not logical_expert_ids.is_contiguous() + or logical_expert_ids.device != self.scheduler.device + ): + raise ValueError("Routes must be contiguous CUDA int32 [tokens, topk] on this device") + num_tokens = int(logical_expert_ids.shape[0]) + if num_tokens > self.max_tokens_per_rank: + raise RuntimeError(f"Route chunk {num_tokens} exceeds {self.max_tokens_per_rank}") + + # Both waits capture their event's current record before a later plan + # can re-record it. No host synchronization or temporary event is needed. + logical_expert_ids.record_stream(self.copy_stream) + self._record_event(self._input_ready_handle, self._execution_stream_handle) + self._wait_event(self._stream_handle, self._input_ready_handle) + outputs = self.scheduler.submit(logical_expert_ids, self._stream_handle) + self._scheduled_generation = self._generation + 1 + self._record_event(self._plan_ready_handle, self._stream_handle) + part = (outputs, num_tokens) + self._plan_part = part + if self._pending_release is not None: + self.lease.prove(self._pending_release[0]) + self._pending_release = None + self._submit_copy(outputs) + return part + + def _submit_copy(self, outputs: object) -> None: + """Consume the HALO-Q plan and enqueue its TMA kernel on COPY.""" + ticket = self.broadcaster.submit(outputs) + if int(ticket.generation) != self._scheduled_generation: + raise RuntimeError("TMA and HALO-Q generation mismatch") + self._generation = int(ticket.generation) + + def plan_finish( + self, part: Tuple[Any, int], *, defer_wait: bool = False + ) -> Tuple[torch.Tensor, int]: + """Lend routes until finish(); optionally defer their device dependency.""" + self._check_owner() + if part is not self._plan_part or part is None: + raise RuntimeError("plan_finish requires this group's pending plan") + outputs, num_tokens = part + if not defer_wait: + self.wait_for_routes() + physical = outputs.physical_slot_ids[:num_tokens] + self.plan_calls += 1 + self._plan_part = None + return physical, self._generation + + def wait_for_routes(self) -> None: + """Order the first MAIN route read after HALO-Q, without a host wait.""" + self._check_owner() + if self._generation <= 0: + raise RuntimeError("No scheduler generation is available") + if self._route_wait_generation != self._generation: + self._wait_event(self._execution_stream_handle, self._plan_ready_handle) + self._route_wait_generation = self._generation + + def plan( + self, logical_expert_ids: torch.Tensor, *, defer_wait: bool = False + ) -> Tuple[torch.Tensor, int]: + return self.plan_finish(self.plan_schedule(logical_expert_ids), defer_wait=defer_wait) + + def discard_plan(self, part: Tuple[Any, int]) -> None: + """Release an enqueued generation that has no MegaMoE consumer.""" + self.plan_finish(part, defer_wait=False) + self.finish() + + def finish(self, completion_event: torch.cuda.Event | None = None) -> None: + """Queue the previous consumer wait before any next-generation producer.""" + self._check_owner() + if ( + self._plan_part is not None + or self._generation <= 0 + or self._generation != self.plan_calls + or self._route_wait_generation != self._generation + or self._generation != self._finished_generation + 1 + ): + raise RuntimeError("finish must pair exactly once with a completed plan") + if completion_event is None: + completion_event = self._consumer_done + self._record_event(self._consumer_done_handle, self._execution_stream_handle) + self.broadcaster.release_generation_after(self._generation, completion_event) + self._pending_release = (self._generation, completion_event) + self._finished_generation = self._generation + + +def build_rebalance_slot_scheduler_group(backend: Any) -> RebalanceSlotSchedulerGroup: + """Build the layer's scheduler/copy group after its live arena exists.""" + arena = getattr(backend, "_rebalance_arena", None) + if arena is None: + raise RuntimeError( + "MoE rebalance requires create_weights to allocate the live arena " + "before building the producer group." + ) + home = getattr(backend, "_rebalance_home_experts", None) + if type(home) is not int or home <= 0: + raise RuntimeError( + f"MoE rebalance requires a positive resident expert count; got {home!r}." + ) + layer_balancer = getattr(backend, "layer_load_balancer", None) + manager = getattr(layer_balancer, "manager", None) + get_copy_stream = getattr(manager, "get_copy_stream", None) + get_ep_comm = getattr(manager, "get_ep_comm", None) + if not callable(get_copy_stream) or not callable(get_ep_comm): + raise RuntimeError( + "per-layer dynamic EPLB scheduler requires model-owned stream and communicator" + ) + return RebalanceSlotSchedulerGroup( + arena=arena, + mapping=backend.mapping, + ep_comm=get_ep_comm(backend._ep_pg), + home_experts=home, + helper_slots=int(backend._rebalance_slots_active), + max_tokens_per_rank=int(backend._maxt_buckets[-1]), + topk=int(backend.routing_method.experts_per_token), + device=int(backend.mapping.local_rank), + copy_stream=get_copy_stream(int(backend.mapping.local_rank)), + auxiliary_sms=int(manager.auxiliary_sms), + layer_idx=getattr(backend, "layer_idx", None), + ) + + +def bind_live_weight_planes(quant_method: Any, module: Any) -> None: + """Alias MegaMoE's seven live parameters onto its rebalance arena.""" + arena = getattr(module, "_rebalance_arena", None) + if arena is None: + return + expected_bundle = quant_method.live_weight_plane_spec(module) + if expected_bundle != arena.bundle: + raise RuntimeError("MoE rebalance arena does not match the quant-method plane spec") + expected_names = tuple(plane.name for plane in expected_bundle.planes) + if tuple(arena.plane_names) != expected_names: + raise RuntimeError("MoE rebalance arena plane order is not canonical") + rebound = False + for name, arena_view, alias in zip( + arena.plane_names, arena.local_plane_views, arena.tekit_alias_views + ): + old = getattr(module, name, None) + if old is None: + raise RuntimeError(f"MoE rebalance live weight plane {name!r} was not registered") + if ( + tuple(alias.shape) != tuple(old.shape) + or tuple(alias.stride()) != tuple(old.stride()) + or alias.dtype != old.dtype + ): + raise RuntimeError(f"MoE rebalance arena alias for {name!r} has an incompatible layout") + if not alias.is_cuda or int(alias.data_ptr()) != int(arena_view.data_ptr()): + raise RuntimeError(f"MoE rebalance arena alias for {name!r} is not its allocator view") + if old.is_meta: + setattr( + module, + name, + torch.nn.Parameter(alias, requires_grad=old.requires_grad), + ) + rebound = True + else: + try: + old.data = alias + except RuntimeError as error: + raise RuntimeError( + f"MoE rebalance failed to bind live plane {name!r}: {error}" + ) from error + if rebound: + quant_method.setup_quant_scales(module) + norm_const = getattr(module, "fc1_norm_const", None) + if norm_const is not None and arena.home_experts < norm_const.data.shape[0]: + norm_const.data[int(arena.home_experts) :].fill_(1.0) + quant_method.assert_live_weight_aliases(module, arena) diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py b/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py index 21943f2f06a4..65113c0fa609 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py @@ -807,8 +807,12 @@ def __init__(self, self.next_layer_repeated_count = None def __del__(self): - if not self.is_shutdown: - self.shutdown() + # MPI collectives are unsafe from rank-asymmetric Python finalization. + # Normal worker teardown calls shutdown() explicitly on every rank. + if not getattr(self, "is_shutdown", True): + logger.warning( + "MoeLoadBalancer was finalized before explicit collective shutdown" + ) def is_static_routing(self): # if we don't update, then it is statistic routing. @@ -949,10 +953,13 @@ def end_iter(self): self.in_iter = False self.iter_id += 1 - def shutdown(self): - """ - Shutdown the load balancer and release resources. - """ + def shutdown(self, local_safe: bool = True) -> bool: + """Collectively release resources after an all-rank safety preflight.""" + if self.is_shutdown: + return True + local_safe = bool(local_safe and not self.in_iter) + if not all(self.shared_mpi_comm.allgather(local_safe)): + return False for single_layer_load_balancer in self.single_layer_load_balancers: single_layer_load_balancer.py_pre_shutdown_cleanup() self.load_balancer_impl.shutdown() @@ -962,6 +969,7 @@ def shutdown(self): single_layer_load_balancer.py_post_shutdown_cleanup() self.shared_mpi_comm.barrier() self.is_shutdown = True + return True def __repr__(self): """ @@ -1050,22 +1058,56 @@ def maybe_create_moe_load_balancer( using_smart_router = mapping and mapping.moe_cluster_size > 1 moe_load_balancer: AbstractContextManager[MoeLoadBalancer | None] = nullcontext() - if in_supported_model_arch and using_ep and not using_smart_router and model_config.moe_load_balancer is not None: - model_config.moe_load_balancer.setup(ep_rank=ep_rank, ep_size=ep_size) - if model_config.moe_load_balancer.layer_updates_per_iter > 0: - # TODO: remove this when supported. - # cpu_arch = platform.machine().lower() - # assert cpu_arch == 'aarch64', "online load balancer only support aarch64, e.g. GB200 now, x86 coming soon." - pass - - moe_load_balancer = MoeLoadBalancer( - ep_rank=ep_rank, - ep_size=ep_size, - layer_updates_per_iter=model_config.moe_load_balancer. - layer_updates_per_iter) - logger.info( - f"Created MoE LoadBalancer, layer_updates_per_iter={model_config.moe_load_balancer.layer_updates_per_iter}..." - ) + config = model_config.moe_load_balancer + mode = getattr(config, "mode", "standard") + if config is not None and mode == "per_layer_dynamic": + if not in_supported_model_arch: + raise ValueError( + "per-layer dynamic EPLB does not support model architecture " + f"{model_arch!r}") + if not using_ep: + raise ValueError("per-layer dynamic EPLB requires expert " + "parallelism (moe_ep_size > 1)") + if using_smart_router: + raise ValueError( + "per-layer dynamic EPLB is incompatible with smart routing " + "(moe_cluster_size > 1)") + + if (in_supported_model_arch and using_ep and not using_smart_router + and config is not None): + config.setup(ep_rank=ep_rank, ep_size=ep_size) + if mode == "per_layer_dynamic": + from .per_layer_dynamic_eplb import PerLayerDynamicMoeLoadBalancer + + moe_load_balancer = PerLayerDynamicMoeLoadBalancer( + ep_rank=ep_rank, + ep_size=ep_size, + config=config, + mapping=mapping, + ) + logger.info( + "Created per-layer dynamic MoE LoadBalancer, num_slots=%s, " + "auxiliary_sms=%s...", + config.num_slots, + getattr(config, "auxiliary_sms", 8), + ) + elif mode == "standard": + if config.layer_updates_per_iter > 0: + # TODO: remove this when supported. + # cpu_arch = platform.machine().lower() + # assert cpu_arch == 'aarch64', "online load balancer only support aarch64, e.g. GB200 now, x86 coming soon." + pass + + moe_load_balancer = MoeLoadBalancer( + ep_rank=ep_rank, + ep_size=ep_size, + layer_updates_per_iter=config.layer_updates_per_iter) + logger.info( + "Created standard MoE LoadBalancer, layer_updates_per_iter=%s...", + config.layer_updates_per_iter, + ) + else: + raise ValueError(f"Unknown MoE load-balancer mode: {mode!r}") return moe_load_balancer diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py b/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py index 932bb57cb130..fa6d6f95bc6a 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py @@ -374,6 +374,7 @@ def build_moe_deployment( # num_experts, so the deployment must say the same. eplb_enabled = get_moe_load_balancer() is not None balancer_config = getattr(model_config, "moe_load_balancer", None) + eplb_mode = getattr(balancer_config, "mode", "standard") if eplb_enabled else "disabled" num_slots = getattr(balancer_config, "num_slots", None) if eplb_enabled else None if num_slots is None: num_slots = num_experts if num_experts is not None else 0 @@ -390,6 +391,7 @@ def build_moe_deployment( env=environment if environment is not None else collect_moe_environment(), # Registered balancer only; config alone does not enable EPLB. eplb_enabled=eplb_enabled, + eplb_mode=eplb_mode, # Routed-expert LoRA only; attention-only LoRA stays False. moe_lora_enabled=has_moe_lora_targets(lora_config), force_dynamic_quantization=model_config.force_dynamic_quantization, @@ -572,6 +574,27 @@ def resolve_moe_impl( rejected = [] eligible: List[MoEImplClass] = [] for candidate in candidates: + if ( + deployment.eplb_mode == "per_layer_dynamic" + and not candidate.capabilities.supports_per_layer_dynamic_eplb + ): + rejected.append( + MoERejection( + _legacy_backend_name(candidate), + MoERejectReason.EPLB_UNSUPPORTED, + f"{candidate.__name__} does not consume per-layer dynamic EPLB plans", + ) + ) + continue + if deployment.eplb_enabled and not candidate.capabilities.supports_eplb: + rejected.append( + MoERejection( + _legacy_backend_name(candidate), + MoERejectReason.EPLB_UNSUPPORTED, + f"{candidate.__name__} does not support EPLB slot layouts", + ) + ) + continue if deployment.moe_lora_enabled and not candidate.capabilities.supports_moe_lora: rejected.append( MoERejection( diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py index 33b9aaf14488..f97969c566e9 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py @@ -1517,6 +1517,8 @@ def _forward_chunk( ) # ----- EPLB: update stats + remap expert ids -> slot ids ----- + uses_replica_plan = moe._load_balancer_uses_replica_plan() + replica_plan = None if moe.layer_load_balancer: moe._load_balancer_done_wait_gpu_stage(is_first_call) # ignore_allreduce=False: the fused kernel has no side channel @@ -1528,7 +1530,13 @@ def _forward_chunk( is_last_call, ignore_allreduce=False, ) - token_selected_slots = moe._load_balancer_route(token_selected_experts, moe.use_dp) + # Standard EPLB remaps here. Per-layer dynamic EPLB must consume the + # final logical routes after calibration/replay below. + token_selected_slots = ( + token_selected_experts + if uses_replica_plan + else moe._load_balancer_route(token_selected_experts, moe.use_dp) + ) else: token_selected_slots = token_selected_experts @@ -1539,42 +1547,51 @@ def _forward_chunk( moe.num_slots, token_selected_slots ) - # ----- quantize / prepare ----- - if getattr(moe.backend, "supports_fused_prepare", lambda: False)(): - # MegaMoE can fuse BF16->MXFP8 quantization with the SymmBuffer - # topk copies, so keep the original activations and let run_moe - # prepare its workspace. - moe_input = x_chunk_real - x_sf = None - else: - # Delegate to ``backend.quantize_input`` so each fused-comm backend - # owns its own empty-tensor layout. Both MegaMoEDeepGemm and - # MegaMoECuteDsl short-circuit ``x.shape[0] == 0`` inside their - # quantize_input contracts. - moe_input, x_sf = moe.backend.quantize_input(x_chunk_real) - - # CuteDSL needs the scheduler's rank-identical chunk maximum to select - # one adaptive bucket on every EP rank; using a local token count could - # diverge and deadlock its in-kernel NVLink barrier. - set_adaptive = getattr(moe.backend, "set_adaptive_launch_tokens", None) - if set_adaptive is not None: - set_adaptive(max(all_rank_num_tokens) if all_rank_num_tokens else None) - - # ----- MoE compute ----- - # ``token_selected_slots`` is in [0, num_slots), matching the kernel's - # ``num_experts`` template parameter (SymmBuffer / weights sized to - # num_slots in quantization.py). - # Fused-comm backends own the EP exchange, so there is no comm plan: - # nothing outside the fused kernel decided anything about this forward. - out = moe.backend.run_moe( - MoERunContext( - token_selected_experts=token_selected_slots, - token_final_scales=token_final_scales, - x=moe_input, - x_sf=x_sf, - output_dtype=output_dtype, + if uses_replica_plan: + token_selected_slots = moe._load_balancer_route(token_selected_slots, moe.use_dp) + replica_plan = moe._load_balancer_replica_plan() + + try: + # ----- quantize / prepare ----- + if getattr(moe.backend, "supports_fused_prepare", lambda: False)(): + # MegaMoE can fuse BF16->MXFP8 quantization with SymmBuffer + # topk copies, so keep activations for run_moe preparation. + moe_input = x_chunk_real + x_sf = None + else: + # Each fused-comm backend owns its empty-tensor layout. + moe_input, x_sf = moe.backend.quantize_input(x_chunk_real) + + # Select one rank-identical adaptive bucket on every EP rank. + set_adaptive = getattr(moe.backend, "set_adaptive_launch_tokens", None) + if set_adaptive is not None: + set_adaptive(max(all_rank_num_tokens) if all_rank_num_tokens else None) + + # ----- MoE compute ----- + # The manager owns the plan; the backend only consumes its + # geometry, READY terminals and late route dependency. + out = moe.backend.run_moe( + MoERunContext( + token_selected_experts=token_selected_slots, + token_final_scales=token_final_scales, + x=moe_input, + x_sf=x_sf, + output_dtype=output_dtype, + replica_plan=replica_plan, + ) ) - ) + except BaseException as error: # noqa: BLE001 + if replica_plan is not None: + try: + moe._load_balancer_finish_replica_plan(replica_plan) + except BaseException as cleanup_error: # noqa: BLE001 + error.add_note( + "Failed to release the per-layer dynamic EPLB plan after " + f"forward failure: {cleanup_error!r}" + ) + raise + else: + moe._load_balancer_finish_replica_plan(replica_plan) # ----- EPLB: start/done CPU rebalance, AFTER run_moe ----- # The external-comm path overlaps CPU stage with ``comm.combine``; diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py new file mode 100644 index 000000000000..ce8b53558a80 --- /dev/null +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py @@ -0,0 +1,564 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Model-owned lifecycle for per-layer dynamic expert load balancing. + +The standard EPLB implementation couples route updates to host weight +migration. Per-layer dynamic EPLB instead produces one device-side replica plan +for every MoE invocation and copies helper weights on a CUDA stream. These +classes preserve the existing load-balancer seams while keeping the scheduler, +arena, and their teardown under one model-level owner. +""" + +from __future__ import annotations + +import weakref +from typing import Any, Callable, Optional, Sequence + +import torch + +from .moe_load_balancer import MoeLoadBalancer + + +class PerLayerDynamicMoeLayerLoadBalancer: + """Adapter between the common MoE routing seam and one replica planner.""" + + def __init__( + self, + manager: "PerLayerDynamicMoeLoadBalancer", + *, + layer_idx: int, + expert_count: int, + top_k: int, + compute_slots_per_rank: int, + repeated_count: int = 1, + ) -> None: + if expert_count <= 0 or expert_count % manager.ep_size != 0: + raise ValueError( + "Per-layer dynamic EPLB requires the expert count to be positive " + f"and divisible by EP size; got E={expert_count}, EP={manager.ep_size}." + ) + resident = expert_count // manager.ep_size + compute = int(compute_slots_per_rank) + if compute <= resident: + raise ValueError( + "Per-layer dynamic EPLB requires at least one helper slot per rank; " + f"got H={resident}, M={compute}." + ) + if top_k <= 0 or repeated_count <= 0: + raise ValueError("top_k and repeated_count must be positive") + + self.manager = manager + self.layer_idx = int(layer_idx) + self.expert_count = int(expert_count) + self.top_k = int(top_k) + self.resident_experts_per_rank = resident + self.compute_slots_per_rank = compute + self.helper_slots_per_rank = compute - resident + self.repeated_count = int(repeated_count) + + self._arena: Optional[Any] = None + self._scheduler_group: Optional[Any] = None + self._bound_backend_ref: Optional[weakref.ReferenceType[Any]] = None + self._pending_replica_plan: Optional[Any] = None + + def get_layer_idx(self) -> int: + return self.layer_idx + + def get_repeat_count(self) -> int: + return self.repeated_count + + def get_load_expert_ids(self) -> list[int]: + start = self.manager.ep_rank * self.resident_experts_per_rank + return list(range(start, start + self.resident_experts_per_rank)) + + def is_static_routing(self) -> bool: + return False + + def is_dynamic_routing(self) -> bool: + """Keep the legacy host-migration path disabled. + + The common route seam still calls :meth:`route` whenever a layer + adapter exists. The legacy meaning of ``is_dynamic_routing`` is + broader: it enables CPU statistics and host-weight migration, neither + of which belongs to the per-layer dynamic CUDA planner. + """ + return False + + def need_load_shared_weights(self) -> bool: + return False + + def set_initial_weight_assignments( + self, initial_weight_assignments: Optional[Sequence[int]] + ) -> None: + if initial_weight_assignments is None: + return + assignments = tuple(int(expert) for expert in initial_weight_assignments) + canonical = tuple(range(self.expert_count)) + if assignments != canonical: + raise ValueError( + "Per-layer dynamic EPLB requires the canonical initial assignment " + "range(num_experts); custom initial_global_assignments are only " + "valid in standard mode." + ) + + # These hooks are part of the standard EPLB layer interface. Routing and + # statistics are produced by HALO-Q, so the corresponding CPU stages are + # intentionally empty here. + def start_wait_gpu_stage(self) -> None: + pass + + def done_wait_gpu_stage(self) -> None: + pass + + def start_set_cpu_stage(self) -> None: + pass + + def done_set_cpu_stage(self) -> None: + pass + + def update_local_statistic(self, *args, **kwargs) -> None: + pass + + def update_statistic_with_local_ids(self, *args, **kwargs) -> None: + pass + + def update_statistic_with_global_ids(self, *args, **kwargs) -> None: + pass + + def update_statistic_with_gathered_statistic(self, *args, **kwargs) -> None: + pass + + def get_local_statistic_tensor(self) -> None: + return None + + def bind_runtime(self, *, arena: Any, scheduler_group: Any) -> None: + """Attach resources built after final weight storage is available.""" + if arena is None or scheduler_group is None: + raise ValueError("arena and scheduler_group are both required") + if self._arena is not None: + if arena is self._arena and scheduler_group is self._scheduler_group: + return + raise RuntimeError(f"per-layer dynamic EPLB layer {self.layer_idx} is already bound") + + arena_home = getattr(arena, "home_experts", self.resident_experts_per_rank) + arena_helpers = getattr(arena, "helper_slots", self.helper_slots_per_rank) + group_home = getattr(scheduler_group, "home_experts", self.resident_experts_per_rank) + group_helpers = getattr(scheduler_group, "helper_slots", self.helper_slots_per_rank) + if (int(arena_home), int(arena_helpers)) != ( + self.resident_experts_per_rank, + self.helper_slots_per_rank, + ): + raise ValueError("arena H/S geometry does not match the EPLB layer") + if (int(group_home), int(group_helpers)) != ( + self.resident_experts_per_rank, + self.helper_slots_per_rank, + ): + raise ValueError("scheduler H/S geometry does not match the EPLB layer") + + self._arena = arena + self._scheduler_group = scheduler_group + self.manager._register_runtime(self) + + def bind_backend(self, backend: Any) -> None: + """Build and own the producer after the backend's live aliases exist.""" + arena = getattr(backend, "_rebalance_arena", None) + if arena is None: + raise RuntimeError("per-layer dynamic EPLB backend has no live arena") + if self._arena is not None or self._scheduler_group is not None: + bound_backend = ( + self._bound_backend_ref() if self._bound_backend_ref is not None else None + ) + if bound_backend is not backend: + raise RuntimeError( + f"per-layer dynamic EPLB layer {self.layer_idx} is already " + "bound to a different backend" + ) + if self._arena is not arena: + raise RuntimeError( + f"per-layer dynamic EPLB layer {self.layer_idx} backend " + "changed its rebalance arena after binding" + ) + if self._scheduler_group is None: + raise RuntimeError( + f"per-layer dynamic EPLB layer {self.layer_idx} has an incomplete runtime binding" + ) + return + + backend_ref = weakref.ref(backend) + from .mega_moe.rebalance_slot_scheduler import build_rebalance_slot_scheduler_group + + if bool(getattr(backend, "tactic_autotune", False)): + # Synthetic profiling may visit every helper row before a real + # plan selects it. Seed those rows from resident weights once; + # serving generations still replace active rows and publish READY. + home = int(arena.home_experts) + with torch.no_grad(): + for plane in arena.tekit_alias_views: + for helper in range(int(arena.helper_slots)): + plane[home + helper].copy_(plane[helper % home], non_blocking=True) + backend._rebalance_autotune_helpers_initialized = True + + scheduler_group = build_rebalance_slot_scheduler_group(backend) + self.bind_runtime(arena=arena, scheduler_group=scheduler_group) + self._bound_backend_ref = backend_ref + + @property + def replica_plan(self) -> Optional[Any]: + return self._pending_replica_plan + + def route( + self, + token_selected_experts: torch.Tensor, + offset_by_ep_rank: bool = False, + *, + gap_hook: Optional[Callable[[], None]] = None, + ) -> torch.Tensor: + """Submit HALO-Q/TMA and return its physical routes without a host wait. + + A shared-expert hook runs after the copy stream is submitted but before + the physical routes are lent to the caller, preserving the intended + scheduler/copy overlap. + """ + del offset_by_ep_rank # HALO-Q consumes global logical expert IDs. + if gap_hook is not None and not callable(gap_hook): + raise TypeError("gap_hook must be callable") + if self._scheduler_group is None or self._arena is None: + raise RuntimeError(f"dynamic EPLB layer {self.layer_idx} has no bound runtime") + if self._pending_replica_plan is not None: + raise RuntimeError(f"dynamic EPLB layer {self.layer_idx} has an unfinished plan") + + if gap_hook is None: + physical, generation = self._scheduler_group.plan( + token_selected_experts, defer_wait=True + ) + else: + part = self._scheduler_group.plan_schedule(token_selected_experts) + try: + gap_hook() + except BaseException as error: # noqa: BLE001 + try: + self._scheduler_group.discard_plan(part) + cleanup_note = "the unused rebalance generation was safely released." + except BaseException as cleanup_error: # noqa: BLE001 + cleanup_note = ( + "rebalance generation cleanup failed and the group " + f"requires teardown: {cleanup_error!r}" + ) + notes = getattr(error, "__notes__", None) + if type(notes) is not list: + notes = [] + error.__notes__ = notes + notes.append("Shared-expert hook failed after HALO-Q/TMA enqueue; " + cleanup_note) + raise + physical, generation = self._scheduler_group.plan_finish(part, defer_wait=True) + generation = int(generation) + self._arena.publish_ready_generation(generation) + + from .impl_contract import MoEReplicaPlan + + plan = MoEReplicaPlan( + resident_experts_per_rank=self.resident_experts_per_rank, + compute_slots_per_rank=self.compute_slots_per_rank, + ready_flags=self._arena.terminal_flags_tensor(), + ready_generation=generation, + reserved_sms=self.manager.auxiliary_sms, + _wait_for_routes=self._scheduler_group.wait_for_routes, + ) + self._pending_replica_plan = plan + return physical + + def require_replica_plan(self) -> Any: + if self._pending_replica_plan is None: + raise RuntimeError(f"dynamic EPLB layer {self.layer_idx} has not produced a plan") + return self._pending_replica_plan + + def finish_replica_plan( + self, + plan: Optional[Any] = None, + *, + completion_event: Optional[torch.cuda.Event] = None, + ) -> None: + pending = self._pending_replica_plan + if pending is None: + raise RuntimeError(f"dynamic EPLB layer {self.layer_idx} has no plan to finish") + if plan is not None and plan is not pending: + raise RuntimeError("attempted to finish a different replica plan") + assert self._scheduler_group is not None + # Also makes cleanup safe when input preparation failed before the + # backend reached its first route consumer. + pending.wait_for_routes() + self._scheduler_group.finish(completion_event=completion_event) + self._pending_replica_plan = None + + @property + def has_submission_owner(self) -> bool: + group = self._scheduler_group + return bool(group is not None and group.has_submission_owner) + + def handoff_submission_owner(self, execution_stream_handle: int) -> None: + """Release this layer's warmup submitter after a caller device drain.""" + group = self._scheduler_group + if group is not None and group.has_submission_owner: + group.release_warmup_owner_after_sync(int(execution_stream_handle)) + + def _clear_runtime(self) -> None: + self._pending_replica_plan = None + self._scheduler_group = None + self._arena = None + self._bound_backend_ref = None + + +class PerLayerDynamicMoeLoadBalancer(MoeLoadBalancer): + """Model-level owner for per-layer dynamic EPLB layers and GPU resources.""" + + def __init__( + self, + *, + ep_rank: int, + ep_size: int, + config: Any, + mapping: Optional[Any] = None, + ep_process_group: Optional[Any] = None, + ) -> None: + # Deliberately do not call MoeLoadBalancer.__init__: per-layer dynamic EPLB + # has no C++ standard-EPLB object and must not create an MPI split. + if ep_size <= 1 or not 0 <= ep_rank < ep_size: + raise ValueError(f"invalid EPLB EP rank/size: rank={ep_rank}, size={ep_size}") + num_slots = int(config.num_slots) + if num_slots <= 0 or num_slots % ep_size != 0: + raise ValueError( + f"num_slots must be positive and divisible by EP size; got {num_slots}" + ) + + self.is_shutdown = False + self.ep_rank = int(ep_rank) + self.ep_size = int(ep_size) + self.config = config + self.mapping = mapping + self.ep_process_group = ep_process_group + self.num_slots = num_slots + self.num_local_slots = num_slots // ep_size + self.auxiliary_sms = int(getattr(config, "auxiliary_sms", 8)) + if self.auxiliary_sms <= 0: + raise ValueError("auxiliary_sms must be positive") + + # Compatibility state used by existing context/engine call sites. + self.layer_updates_per_iter = 1 + self._previous_balancer = None + self.single_layer_load_balancers: list[PerLayerDynamicMoeLayerLoadBalancer] = [] + self.next_layer_repeated_count: Optional[int] = None + self.iter_id = 0 + self.in_iter = False + self.enable_statistic = False + self.enable_update_weights = False + self._copy_stream: Optional[torch.cuda.Stream] = None + self._copy_stream_device: Optional[int] = None + self._ep_comm: Optional[Any] = None + + # Providers use this manager as their scope owner. Resource references + # stay reachable here, never through Mapping or a process-global registry. + self._arena_providers: list[Any] = [] + self._rebalance_shared_slot_pools: dict[tuple, Any] = {} + self._rebalance_shared_slot_pool_order: list[Any] = [] + + def __del__(self) -> None: + # Teardown is collective and must be driven explicitly by the worker; + # a Python finalizer cannot safely enter EP collectives. + pass + + def is_static_routing(self) -> bool: + return False + + def is_dynamic_routing(self) -> bool: + return True + + def set_repeated_for_next_layer(self, repeated_count: int) -> None: + if repeated_count <= 0: + raise ValueError("repeat count must be positive") + self.next_layer_repeated_count = int(repeated_count) + + def add_layer( + self, + expert_count: int, + top_k: int, + slot_count_per_rank: int, + aux_stream: Optional[torch.cuda.Stream] = None, + ) -> PerLayerDynamicMoeLayerLoadBalancer: + del aux_stream # The scheduler owns its high-priority copy stream. + self.config.validate_expert_capacity(expert_count) + if int(slot_count_per_rank) != self.num_local_slots: + raise ValueError( + "per-layer dynamic EPLB layer capacity must match the configured " + f"local slot count; got {slot_count_per_rank}, expected " + f"{self.num_local_slots}" + ) + repeated_count = self.next_layer_repeated_count or 1 + self.next_layer_repeated_count = None + layer = PerLayerDynamicMoeLayerLoadBalancer( + self, + layer_idx=len(self.single_layer_load_balancers), + expert_count=expert_count, + top_k=top_k, + compute_slots_per_rank=slot_count_per_rank, + repeated_count=repeated_count, + ) + self.single_layer_load_balancers.append(layer) + return layer + + def get_copy_stream(self, device: int) -> torch.cuda.Stream: + """Return this model's single highest-priority auxiliary stream.""" + device = int(device) + if self._copy_stream is None: + total_sms = int(torch.cuda.get_device_properties(device).multi_processor_count) + if self.auxiliary_sms >= total_sms: + raise ValueError( + "auxiliary_sms must leave at least one SM for the compute stream; " + f"got {self.auxiliary_sms} of {total_sms}" + ) + _, priority = torch.cuda.Stream.priority_range() + self._copy_stream = torch.cuda.Stream(device=device, priority=priority) + self._copy_stream_device = device + elif self._copy_stream_device != device: + raise RuntimeError("one per-layer dynamic EPLB manager cannot span CUDA devices") + return self._copy_stream + + def get_ep_comm(self, process_group: Any) -> Any: + """Return this model's adapter for its resolved EP ProcessGroup.""" + from .mega_moe.rebalance_live_arena import _EpComm + + if self._ep_comm is None: + self._ep_comm = _EpComm(process_group) + elif self._ep_comm.process_group is not process_group: + raise RuntimeError("one per-layer dynamic EPLB manager cannot span EP ProcessGroups") + self._ep_comm.validate_geometry(expected_rank=self.ep_rank, expected_world=self.ep_size) + return self._ep_comm + + def create_arena_provider(self) -> Any: + """Create a per-layer provider scoped to this model manager.""" + from .mega_moe.rebalance_live_arena import SharedSlotArenaProvider + + provider = SharedSlotArenaProvider(self) + self._arena_providers.append(provider) + return provider + + def adopt_arena_provider(self, provider: Any) -> None: + """Retain a transitional provider already created by the allocator.""" + if provider not in self._arena_providers: + self._arena_providers.append(provider) + + def _register_runtime(self, layer: PerLayerDynamicMoeLayerLoadBalancer) -> None: + if layer not in self.single_layer_load_balancers: + raise ValueError("cannot bind a layer owned by another EPLB manager") + assert layer._arena is not None + self.adopt_arena_provider(layer._arena.provider) + + def set_iter_info( + self, + enable_statistic: Optional[bool], + enable_update_weights: Optional[bool], + ) -> None: + if enable_statistic is not None: + self.enable_statistic = bool(enable_statistic) + if enable_update_weights is not None: + self.enable_update_weights = bool(enable_update_weights) + + def start_iter(self) -> None: + if self.in_iter: + raise RuntimeError("already in a MoE load-balancer iteration") + self.in_iter = True + + def end_iter(self) -> None: + if not self.in_iter: + raise RuntimeError("not in a MoE load-balancer iteration") + unfinished = [ + layer.layer_idx + for layer in self.single_layer_load_balancers + if layer.replica_plan is not None + ] + if unfinished: + raise RuntimeError(f"unfinished replica plans at iteration end: {unfinished}") + self.in_iter = False + self.iter_id += 1 + + def set_warmup(self, value: bool) -> None: + """Keep production HALO-Q/TMA planning active during warmup.""" + del value + + def register_weight_slots_after_to_cuda(self) -> None: + # Live aliases are registered by the per-layer dynamic arena binder. + pass + + def finalize_model(self) -> None: + # Scheduler groups are built when each layer's final aliases exist. + pass + + def set_warm_up_iter_count(self, iter_count: int) -> None: + del iter_count + + @property + def has_submission_owner(self) -> bool: + return any(layer.has_submission_owner for layer in self.single_layer_load_balancers) + + def handoff_submission_owner(self, execution_stream_handle: int) -> None: + """Transfer all warmup-owned submitters after one caller device drain.""" + for layer in self.single_layer_load_balancers: + layer.handoff_submission_owner(execution_stream_handle) + + def release_quiesced_submission_owner(self, execution_stream_handle: int) -> None: + """Release executor ownership after the caller drains the device.""" + for layer in self.single_layer_load_balancers: + group = layer._scheduler_group + if group is not None: + group.release_quiesced_submission_owner(int(execution_stream_handle)) + + def shutdown(self, local_safe: bool = True) -> bool: + """Collectively close this model's pools in deterministic layer order.""" + if self.is_shutdown: + return True + + # Agree on local readiness before any rank starts destructive teardown. + # This also covers managers that failed before their first pool existed. + local_safe = bool(local_safe and not self.in_iter) + comm = self._ep_comm + if comm is None and self.ep_process_group is not None: + comm = self.get_ep_comm(self.ep_process_group) + if comm is not None: + states = comm.allgather( + { + "local_safe": local_safe, + "pool_count": len(self._rebalance_shared_slot_pool_order), + } + ) + pool_counts = {state["pool_count"] for state in states} + if len(pool_counts) != 1: + raise RuntimeError( + "per-layer dynamic EPLB pool counts differ across the EP communicator" + ) + local_safe = all(state["local_safe"] for state in states) + if not local_safe: + return False + + # Snapshot creation order because each successful close removes itself. + for pool in list(self._rebalance_shared_slot_pool_order): + if not pool.close_collectively(local_safe=True): + return False + + # Pool teardown closes its registered producers. Calling close again is + # deliberately allowed and also covers a transitional group that has + # not yet been registered with a pool. + for layer in self.single_layer_load_balancers: + group = layer._scheduler_group + if group is not None: + group.close() + layer._clear_runtime() + self._arena_providers.clear() + self._ep_comm = None + self._rebalance_shared_slot_pools.clear() + self._rebalance_shared_slot_pool_order.clear() + self._copy_stream = None + self._copy_stream_device = None + self.is_shutdown = True + return True + + +__all__ = ["PerLayerDynamicMoeLayerLoadBalancer", "PerLayerDynamicMoeLoadBalancer"] diff --git a/tensorrt_llm/_torch/moe/fused_moe/quantization.py b/tensorrt_llm/_torch/moe/fused_moe/quantization.py index a5a850409f59..fd95c03dffc6 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/moe/fused_moe/quantization.py @@ -4282,38 +4282,46 @@ def _get_fc2_alpha_input_scale( return entry['w2'][...].reshape([]).to(dtype=torch.float32).reciprocal() # ----------------------------------------------------------------- - # Shape helpers (kernel-side authoritative; the SF flat sizes match - # kernel_fc12.py as cited per-method below). + # Shape helpers delegate to the seven-plane ABI used by the arena. # ----------------------------------------------------------------- @staticmethod - def _ceil_div_int(a: int, b: int) -> int: - return (a + b - 1) // b + def _live_weight_bundle_layout(hidden: int, + intermediate: int, + expand_intermediate: Optional[int] = None): + from ...cute_dsl_kernels.megamoe_scheduler.sami.geometry import \ + BundleLayout + + return BundleLayout.create( + hidden=int(hidden), + intermediate=int(intermediate), + expand_intermediate=(None if expand_intermediate is None else + int(expand_intermediate)), + ) @staticmethod - def _round_up_int(a: int, b: int) -> int: - return ((a + b - 1) // b) * b + def _live_weight_plane(bundle, name: str): + try: + return next(plane for plane in bundle.planes if plane.name == name) + except StopIteration as error: + raise RuntimeError( + f"canonical live-weight bundle is missing plane {name!r}" + ) from error @classmethod def fc1_sf_flat_size(cls, intermediate: int, hidden: int) -> int: - """``round_up(expand_intermediate, SfPaddingBlock=128) * - round_up(ceil(hidden / 16), 4)`` -- matches the FC1 weight-SF view - in ``kernel_fc12.py`` (``intermediate_gateup_padded`` / - ``expected_fc1_weight_sf_cols``). - ``expand_intermediate = 2 * intermediate``. - """ - expand_intermediate = intermediate * 2 - return (cls._round_up_int(expand_intermediate, 128) * - cls._round_up_int(cls._ceil_div_int(hidden, 16), 4)) + """Return the canonical FC1 scale-factor elements per slot.""" + bundle = cls._live_weight_bundle_layout(hidden, intermediate) + plane = cls._live_weight_plane(bundle, "mega_fc1_weight_sf") + assert plane.storage is not None + return plane.storage.elements_per_slot @classmethod def fc2_sf_flat_size(cls, hidden: int, intermediate: int) -> int: - """``round_up(hidden, SfPaddingBlock=128) * - round_up(ceil(intermediate / 16), 4)`` -- matches the FC2 weight-SF - view in ``kernel_fc12.py`` (``hidden_padded_fc2`` / - ``expected_fc2_weight_sf_cols``). - """ - return (cls._round_up_int(hidden, 128) * - cls._round_up_int(cls._ceil_div_int(intermediate, 16), 4)) + """Return the canonical FC2 scale-factor elements per slot.""" + bundle = cls._live_weight_bundle_layout(hidden, intermediate) + plane = cls._live_weight_plane(bundle, "mega_fc2_weight_sf") + assert plane.storage is not None + return plane.storage.elements_per_slot # Source-checkpoint tensors kept as 0-element placeholders outside the # load window so the full source set and the mega buffers never coexist @@ -4327,28 +4335,61 @@ def fc2_sf_flat_size(cls, hidden: int, intermediate: int) -> int: # "peak = steady set + ONE layer" bound. Class-level on purpose. _streamed_transient_lock = threading.Lock() + def live_weight_plane_spec(self, module: torch.nn.Module): + """Return the canonical seven-plane ABI for this quantized module.""" + return self._live_weight_bundle_layout( + hidden=int(module.hidden_size), + intermediate=int(module.intermediate_size_per_partition), + expand_intermediate=int( + module.expand_intermediate_size_per_partition), + ) + + def assert_live_weight_aliases(self, module: torch.nn.Module, + arena) -> None: + """Fail if loading or transformation replaced an arena-backed plane.""" + if self.live_weight_plane_spec(module) != arena.bundle: + raise RuntimeError( + "MoE rebalance arena does not match the quant-method plane spec" + ) + if not (len(arena.bundle.planes) == len(arena.local_plane_views) == len( + arena.tekit_alias_views)): + raise RuntimeError( + "MoE rebalance live-plane collections differ in length") + for plane, arena_view, alias in zip( + arena.bundle.planes, + arena.local_plane_views, + arena.tekit_alias_views, + ): + current = getattr(module, plane.name, None) + if current is None or bool(getattr(current, "is_meta", False)): + raise RuntimeError( + f"MoE rebalance live weight plane {plane.name!r} is not materialized" + ) + if (tuple(current.shape) != tuple(alias.shape) + or tuple(current.stride()) != tuple(alias.stride()) + or current.dtype != alias.dtype + or current.device != alias.device): + raise RuntimeError( + f"MoE rebalance live weight plane {plane.name!r} changed layout after binding" + ) + pointers = ( + int(current.data_ptr()), + int(alias.data_ptr()), + int(arena_view.data_ptr()), + ) + if pointers[0] != pointers[1] or pointers[1] != pointers[2]: + raise RuntimeError( + f"MoE rebalance live weight plane {plane.name!r} changed storage after binding" + ) + # ----------------------------------------------------------------- # create_weights: register MegaMoE-format parameters in addition to # the grandparent's standard NVFP4 parameters. # ----------------------------------------------------------------- def create_weights(self, module: torch.nn.Module): - # The MegaMoE NVFP4 weight + SF pipeline hard-codes the gated - # 2x expansion (``expand_intermediate == 2 * intermediate``): - # ``fc1_sf_flat_size`` computes ``round_up(2 * intermediate, 128)`` - # and ``_build_mega_format_buffers`` slices ``w3_w1_weight`` at - # ``[:intermediate, :]`` / ``[intermediate:, :]`` before the - # 16-atom gate/up interleave. A non-2x configuration would - # silently mis-size the registered ``mega_fc1_weight*`` buffers - # and the loader. Fail fast at create time instead of breaking - # inside ``_build_mega_format_buffers``. - if (module.expand_intermediate_size_per_partition - != 2 * module.intermediate_size_per_partition): - raise NotImplementedError( - "NVFP4MegaMoECuteDslMethod currently requires the gated " - "2x expansion (expand_intermediate == 2 * intermediate); " - f"got expand_intermediate=" - f"{module.expand_intermediate_size_per_partition}, " - f"intermediate={module.intermediate_size_per_partition}.") + # Resolve the same descriptor used by build-time arena allocation before + # the parent materializes any source or derived Parameters. + bundle = self.live_weight_plane_spec(module) # Fail fast: the derived-only EPLB migration path registers no bias # staging. if module.bias and self.need_load_shared_weights(module): @@ -4385,69 +4426,73 @@ def create_weights(self, module: torch.nn.Module): # init tensors are actually released. self.setup_quant_scales(module) - num_local_slots = module.expert_size_per_partition - hidden = module.hidden_size - intermediate = module.intermediate_size_per_partition - expand_intermediate = module.expand_intermediate_size_per_partition - # NVFP4 packs 2 elements per byte along K (= hidden for fc1, = - # intermediate for fc2), so the natural HF ``(slots, N, K_bytes)`` - # storage already has K as the stride-1 (innermost) dim. The MegaMoE - # CuteDSL kernel reads the weight K-major with K innermost; the - # backend hands it a ``.transpose(1, 2)`` VIEW (NOT ``.contiguous()``) - # of this storage so the logical shape becomes ``(slots, K_bytes, N)`` - # while K stays stride-1 (see ``mega_moe_cute_dsl`` kernel-input prep). - # ``mega_fc1_weight`` keeps the 16-atom gate/up interleave along - # expand_intermediate (the N axis). ``mega_fc2_weight`` is byte- - # equivalent to ``w2_weight``. - mega_fc1_weight = nn.Parameter( - torch.empty(num_local_slots, - expand_intermediate, - hidden // 2, - dtype=torch.uint8), - requires_grad=False, - ) - module.register_parameter("mega_fc1_weight", mega_fc1_weight) - - mega_fc2_weight = nn.Parameter( - torch.empty(num_local_slots, - hidden, - intermediate // 2, - dtype=torch.uint8), - requires_grad=False, - ) - module.register_parameter("mega_fc2_weight", mega_fc2_weight) + # The generic NVFP4 source/checkpoint Parameters created above use + # the resident axis H (``expert_size_per_partition``). Only the seven + # tensors consumed directly by the MegaMoE kernel use M = H + S. + num_resident_experts = int(module.expert_size_per_partition) + num_compute_slots = int(module.compute_slot_count) + if num_compute_slots < num_resident_experts: + raise ValueError( + "MegaMoE compute_slot_count must cover every resident expert: " + f"M={num_compute_slots}, H={num_resident_experts}.") + + planes = {plane.name: plane for plane in bundle.planes} + + def storage_metadata(name: str, slots: int): + plane = planes[name] + return plane.storage_metadata(slots) + + def storage_parameter(name: str, factory=torch.empty) -> nn.Parameter: + shape, stride, dtype_name, _, _ = storage_metadata( + name, num_compute_slots) + tensor = factory(shape, dtype=getattr(torch, dtype_name)) + if tuple(tensor.stride()) != stride: + raise RuntimeError( + f"{name} canonical storage descriptor is not contiguous") + return nn.Parameter(tensor, requires_grad=False) + + # The NVFP4 parent owns alpha construction. Derive the resident and + # compute shapes from its canonical planes, preserve the resident + # values, and refresh setup_quant_scales after the final identities are + # installed. + for name in ("fc31_alpha", "fc2_alpha"): + resident = getattr(module, name) + resident_shape, _, dtype_name, _, _ = storage_metadata( + name, num_resident_experts) + target_shape, _, _, _, _ = storage_metadata(name, num_compute_slots) + if (tuple(resident.shape) != resident_shape + or resident.dtype != getattr(torch, dtype_name)): + raise RuntimeError( + f"{name} source layout must be {resident_shape} " + f"{dtype_name}; got {tuple(resident.shape)} " + f"{resident.dtype}.") + if target_shape != resident_shape: + expanded = resident.detach().new_ones(target_shape) + expanded[:num_resident_experts].copy_(resident.detach()) + module.register_parameter( + name, nn.Parameter(expanded, requires_grad=False)) - mega_fc1_weight_sf = nn.Parameter( - torch.empty(num_local_slots, - self.fc1_sf_flat_size(intermediate, hidden), - dtype=torch.uint8), - requires_grad=False, - ) - module.register_parameter("mega_fc1_weight_sf", mega_fc1_weight_sf) + # Only the four derived weight/SF buffers are newly registered here. + # Their storage descriptors match the loader; the launch path exposes + # the two weight buffers as zero-copy transposed kernel views. + for name in ( + "mega_fc1_weight", + "mega_fc1_weight_sf", + "mega_fc2_weight", + "mega_fc2_weight_sf", + ): + module.register_parameter(name, storage_parameter(name)) - mega_fc2_weight_sf = nn.Parameter( - torch.empty(num_local_slots, - self.fc2_sf_flat_size(hidden, intermediate), - dtype=torch.uint8), - requires_grad=False, - ) - module.register_parameter("mega_fc2_weight_sf", mega_fc2_weight_sf) - - # Per-expert FC1-output (= FC2-input) NVFP4 quantization norm_const. - # The MegaMoE CuteDSL kernel ABI is per-expert ``(num_local_slots,)``. - # This buffer is filled in ``process_weights_after_loading`` from each - # local expert's raw ``w2.input_scale`` as ``1 / w2.input_scale`` and is - # a stable, contiguous, device-local tensor (NOT a stride-0 expand view) - # so the runner's - # ``from_dlpack(...).mark_layout_dynamic(...)`` and the compile cache - # see a normal 1-D fp32 layout. Because the value is genuinely - # per-expert, EPLB shared-load paths also register CPU staging for this - # parameter. - fc1_norm_const = nn.Parameter( - torch.ones(num_local_slots, dtype=torch.float32), - requires_grad=False, - ) - module.register_parameter("fc1_norm_const", fc1_norm_const) + # Per-expert FC1-output (= FC2-input) NVFP4 quantization norm_const is + # initialized to one and populated after loading from w2.input_scale. + # The ones come from the factory: MetaInitMode keeps torch.ones tensors + # real but rejects any later write into a meta tensor. + module.register_parameter( + "fc1_norm_const", storage_parameter("fc1_norm_const", torch.ones)) + + # Alpha Parameters may have been resized or rebound after the parent + # built its view; refresh it against the final seven-plane identities. + self.setup_quant_scales(module) def _materialize_source_params(self, module: torch.nn.Module): """Rematerialize this module's streamed source params (full shape) @@ -4701,7 +4746,7 @@ def _rows_covered(set_name: str, param) -> int: def _check_initial_aux_scale_coverage(self, module: torch.nn.Module) -> None: """Reject partially populated NVFP4 auxiliary-scale families.""" - n_slots = module.expert_size_per_partition + n_slots = int(module.expert_size_per_partition) # A whole-checkpoint load is handed every expert's input_scale, because # the weights dict holds the entire checkpoint; a streaming EP load # only ever reads its own rank's experts, so its complete answer is @@ -4771,7 +4816,7 @@ def process_weights_after_loading(self, module: torch.nn.Module): if module.w3_w1_weight.data.numel() == 0: return - n_slots = module.expert_size_per_partition + n_slots = int(module.expert_size_per_partition) coverage = self._streamed_coverage(module) incomplete = {k: v for k, v in coverage.items() if v < n_slots} if incomplete: @@ -4918,7 +4963,12 @@ def _build_fc1_norm_const(self, module: torch.nn.Module) -> None: raw_input_scales, module.initial_local_expert_ids, device=module.fc1_norm_const.device) - module.fc1_norm_const.data.copy_(routed_norm_const) + n_home = int(module.expert_size_per_partition) + if routed_norm_const.numel() != n_home: + raise RuntimeError( + "MegaMoE-CuteDSL fc1_norm_const expected one entry per " + f"resident expert ({n_home}), got {routed_norm_const.numel()}.") + module.fc1_norm_const.data[:n_home].copy_(routed_norm_const) if self.need_load_shared_weights(module): local_shared_load_expert_ids = module.layer_load_balancer.get_load_expert_ids( @@ -5014,11 +5064,24 @@ def _build_mega_format_buffers( passes ``module.w3_w1_weight.data`` etc., the EPLB staging path passes ``module.local_shared_w3_w1_tensors`` etc. """ + bundle = self._live_weight_bundle_layout(hidden, intermediate, + expand_intermediate) + planes = {plane.name: plane for plane in bundle.planes} + + def storage(name: str): + view = planes[name].storage + if view is None: + raise RuntimeError( + f"canonical live-weight plane {name!r} has no storage view") + return view + if intermediate % 16 != 0: raise ValueError( f"MegaMoE NVFP4 FC1 transform requires intermediate % 16 == 0" f" (Fc1GateUpInterleave); got intermediate={intermediate}.") - h_bytes = hidden // 2 + expand_intermediate = bundle.expand_intermediate + assert expand_intermediate is not None + h_bytes = storage("mega_fc1_weight").shape[1] n_pairs = intermediate // 16 # The parent NVFP4 layout stores raw weight tensors as int64 @@ -5046,13 +5109,18 @@ def _build_mega_format_buffers( # 2)`` view ``(slots, hidden//2, expand_intermediate)`` so the kernel # sees K-major with K stride-1 -- WITHOUT materializing a contiguous # copy (which would move K off the innermost axis and corrupt the GEMM). - mega_fc1 = interleaved.view(num_slots, expand_intermediate, - h_bytes).contiguous() + fc1_storage_shape = (num_slots, *storage("mega_fc1_weight").shape) + mega_fc1 = interleaved.view(fc1_storage_shape).contiguous() # ----- FC2 weight: byte-equivalent clone ----- # ``raw_w2`` is ``(slots, hidden, intermediate//2)`` (N, K_bytes) with # K = intermediate//2 innermost; the kernel-input prep transposes the # last two dims to expose K-major as a view (see fc1 note above). + fc2_storage_shape = (num_slots, *storage("mega_fc2_weight").shape) + if tuple(raw_w2.shape) != fc2_storage_shape: + raise ValueError( + "raw FC2 weight does not match canonical storage shape: " + f"got {tuple(raw_w2.shape)}, expected {fc2_storage_shape}") mega_fc2 = raw_w2.detach().clone().contiguous() # ----- FC1 weight SF: same 16-atom interleave + to_blocked ----- @@ -5062,7 +5130,7 @@ def _build_mega_format_buffers( gate_up_interleave_intermediate=intermediate, n_pairs=n_pairs, expand_intermediate=expand_intermediate, - flat_size=self.fc1_sf_flat_size(intermediate, hidden), + flat_size=storage("mega_fc1_weight_sf").elements_per_slot, ) # ----- FC2 weight SF: per-slot to_blocked only (no gate/up) ----- @@ -5072,7 +5140,7 @@ def _build_mega_format_buffers( gate_up_interleave_intermediate=None, n_pairs=None, expand_intermediate=expand_intermediate, - flat_size=self.fc2_sf_flat_size(hidden, intermediate), + flat_size=storage("mega_fc2_weight_sf").elements_per_slot, ) return mega_fc1, mega_fc1_sf, mega_fc2, mega_fc2_sf diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/VENDOR_STAMP b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/VENDOR_STAMP new file mode 100644 index 000000000000..3fbb6a2e0395 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/VENDOR_STAMP @@ -0,0 +1,7 @@ +upstream: https://gitlab-master.nvidia.com/bangyus/cutedsl_megamoe.git +locked_commit: 382ce4a6f111d0769f31c4c975e1c643f0dc4f70 +project: CuTe DSL MegaMoE +original_authors: Bangyu Shen ; Aragorn Guan ; Hao Sheng ; Yuhan Li ; Yucheng Yan ; XingFei Xi ; Shawn Wang +license: per-file SPDX declarations (BSD-3-Clause where declared) +downstream_patch: 3rdparty/vendor_patches/cutedsl-megamoe.patch (TensorRT-LLM main adaptations: SwiGLU alpha/beta for MiniMax-M3, bandit B105 annotations) +managed_by: 3rdparty/vendor_sources.lock.yaml diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/__init__.py index 786802035913..8feb6089feb0 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/__init__.py @@ -1,24 +1 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Exported MegaMoE KernelClass entry points.""" - -from .kernel_src.blackwell.inference.mega.block_scaled_swap_ab_mega_moe_kernel import ( - BlockScaledSwapAbMegaMoeKernel as BlackwellInferenceMegaMoE, -) -from .kernel_src.rubin.inference.local_mega.block_scaled_swap_ab_local_mega_moe_kernel import ( - BlockScaledSwapAbLocalMegaMoeKernel as RubinInferenceLocalMegaMoE, -) -from .kernel_src.rubin.inference.mega.block_scaled_swap_ab_mega_moe_kernel import ( - BlockScaledSwapAbMegaMoeKernel as RubinInferenceMegaMoE, -) -from .kernel_src.rubin.inference.mega.block_scaled_swap_ab_mega_moe_kernel_gen_specialized import ( - BlockScaledSwapAbGenphaseMoeKernel as RubinInferenceGenphaseMegaMoE, -) - -__all__ = [ - "BlackwellInferenceMegaMoE", - "RubinInferenceGenphaseMegaMoE", - "RubinInferenceLocalMegaMoE", - "RubinInferenceMegaMoE", -] +"""Integration-ready CuTeDSL MegaMoE sources.""" diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/api.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/api.py index 06bc3a8d0426..c3fd0af463b0 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/api.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/api.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Stable construction API for composable kernel implementations.""" import types @@ -11,6 +8,7 @@ from cutlass.cute.typing import SymInt + RuntimeIntegerType = SymInt StaticIntegerType = int StaticOrRuntimeIntegerType = Union[int, SymInt] @@ -29,15 +27,16 @@ def __init__(self, expected_type) -> None: def _required_type(requirement: Requirement): - return ( - requirement.expected_type if isinstance(requirement, OptionalRequirement) else requirement - ) + return requirement.expected_type if isinstance(requirement, OptionalRequirement) else requirement def _matches_type(value, expected_type) -> bool: origin = get_origin(expected_type) if origin in (Union, types.UnionType): - return any(_matches_type(value, candidate) for candidate in get_args(expected_type)) + return any( + _matches_type(value, candidate) + for candidate in get_args(expected_type) + ) return isinstance(value, expected_type) @@ -66,7 +65,9 @@ def validate( if name not in self._values: if isinstance(requirement, OptionalRequirement): continue - raise KeyError(f"{component_name} requires descriptor field {name!r}.") + raise KeyError( + f"{component_name} requires descriptor field {name!r}." + ) expected_type = _required_type(requirement) value = self._values[name] if not _matches_type(value, expected_type): @@ -96,11 +97,13 @@ class KernelComponent(ABC, metaclass=_KernelComponentMeta): @classmethod @abstractmethod - def problem_desc_require(cls) -> Dict[str, Requirement]: ... + def problem_desc_require(cls) -> Dict[str, Requirement]: + ... @classmethod @abstractmethod - def impl_desc_require(cls) -> Dict[str, Requirement]: ... + def impl_desc_require(cls) -> Dict[str, Requirement]: + ... def _validate_desc_inputs( self, @@ -108,10 +111,14 @@ def _validate_desc_inputs( impl_desc: ImplDesc, ) -> None: component_name = type(self).__name__ - overlap = self.problem_desc_require().keys() & self.impl_desc_require().keys() + overlap = ( + self.problem_desc_require().keys() + & self.impl_desc_require().keys() + ) if overlap: raise ValueError( - f"{component_name} requires fields from both descriptors: {sorted(overlap)}." + f"{component_name} requires fields from both descriptors: " + f"{sorted(overlap)}." ) problem_desc.validate( self.problem_desc_require(), @@ -132,26 +139,33 @@ def _validate_component_init(self) -> None: if not hasattr(self, name): if isinstance(requirement, OptionalRequirement): continue - raise RuntimeError(f"{component_name} did not bind required field {name!r}.") + raise RuntimeError( + f"{component_name} did not bind required field {name!r}." + ) expected_type = _required_type(requirement) value = getattr(self, name) if not _matches_type(value, expected_type): raise TypeError( - f"{component_name}.{name} must have type {expected_type}, got {type(value)}." + f"{component_name}.{name} must have type " + f"{expected_type}, got {type(value)}." ) for name, value in vars(self).items(): if isinstance(value, Desc): - raise RuntimeError(f"{component_name}.{name} retains a descriptor.") + raise RuntimeError( + f"{component_name}.{name} retains a descriptor." + ) class KernelClass(KernelComponent): """Top-level host wrapper for one composable kernel implementation.""" @abstractmethod - def name(self) -> str: ... + def name(self) -> str: + ... @abstractmethod - def aot_compile(self): ... + def aot_compile(self): + ... __all__ = [ diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/__init__.py index 52a7a9daf028..0e7d0e5264dc 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/__init__.py @@ -1,2 +1,47 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Low-level workspace, synchronization, and PTX helpers.""" + +from .device_workspace import DeviceWorkspace +from .dsl_helpers import spin_peek, spin_wait +from .flag_batch import GpuAsyncReleaseFlagBatchTracker, GpuReleaseFlagBatchTracker, make_flag_batch_tracker +from .iket_compat import iket +from .ptx_helpers import ( + cvt_f32_to_fp8_to_f32, + cvt_f32x4_to_f8x4_pack_i32, + stg_e8m0_from_f32, + stg_e8m0x8_from_f32, +) +from .smem_workspace import SmemWorkspace +from .software_sync import NvlinkBarrier, SoftwareGridSync +from .utils import ( + IntegerType, + ceil_div, + cosize_from_shape_stride_tuples, + product, + row_major_stride, + round_up, + validate_static_integer_tuple, +) + +__all__ = [ + "DeviceWorkspace", + "GpuAsyncReleaseFlagBatchTracker", + "GpuReleaseFlagBatchTracker", + "IntegerType", + "NvlinkBarrier", + "SmemWorkspace", + "SoftwareGridSync", + "ceil_div", + "cosize_from_shape_stride_tuples", + "cvt_f32_to_fp8_to_f32", + "cvt_f32x4_to_f8x4_pack_i32", + "make_flag_batch_tracker", + "product", + "row_major_stride", + "round_up", + "iket", + "spin_peek", + "spin_wait", + "stg_e8m0_from_f32", + "stg_e8m0x8_from_f32", + "validate_static_integer_tuple", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/__init__.py index 52a7a9daf028..337547020a79 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/__init__.py @@ -1,2 +1,26 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Kernel source modules independent of repository-only runners.""" + +from .schedulers import ( + BlackwellFusedFc12Scheduler, + BlockPhase, + Fc12WorkTileState, + NonSwapAbFc12WorkTileInfo, + SchedulerBase, + SchedulerConsumer, + SchedulerWorkTileBase, + SwapAbFc12WorkTileInfo, + WorkIdAcquisitionMode, +) + + +__all__ = [ + "BlackwellFusedFc12Scheduler", + "BlockPhase", + "Fc12WorkTileState", + "NonSwapAbFc12WorkTileInfo", + "SchedulerBase", + "SchedulerConsumer", + "SchedulerWorkTileBase", + "SwapAbFc12WorkTileInfo", + "WorkIdAcquisitionMode", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/symmetric_buffer.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/symmetric_buffer.py index 4ac013c3b5f0..42c8a7b6c20e 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/symmetric_buffer.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/symmetric_buffer.py @@ -1,11 +1,10 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Symmetric-heap peer pointer mapping carried in kernel arguments.""" from dataclasses import dataclass from typing import Any, Optional +from packaging.version import Version + import cutlass import cutlass.cute as cute from cutlass._mlir import ir @@ -15,7 +14,7 @@ from cutlass.base_dsl.typing import get_c_pointers from cutlass.cute.typing import AddressSpace from cutlass.cutlass_dsl import Int32, Int64, dsl_user_op -from packaging.version import Version + try: from cutlass.base_dsl.typing import MLIR_DYNAMIC_INDEX @@ -64,9 +63,7 @@ def __extract_mlir_attributes__(self) -> list: return [ir.DictAttr.get({})] @cute.jit - def map( - self, local_address: Int64, destination_rank: Int32, byte_offset: Int64 = Int64(0) - ) -> Int64: + def map(self, local_address: Int64, destination_rank: Int32, byte_offset: Int64 = Int64(0)) -> Int64: if cutlass.const_expr(self.max_ranks <= _byval_rank_limit): i64_type = ir.Type.parse("i64") offset_pointer = llvm.getelementptr( @@ -89,10 +86,7 @@ def map_pointer(self, pointer, destination_rank: Int32, byte_alignment: Optional if cutlass.const_expr(byte_alignment is None): byte_alignment = pointer.max_alignment return cute.make_ptr( - pointer.dtype, - self.map(pointer.toint(), destination_rank), - pointer.memspace, - assumed_align=byte_alignment, + pointer.dtype, self.map(pointer.toint(), destination_rank), pointer.memspace, assumed_align=byte_alignment ) @@ -119,27 +113,11 @@ def make_device_object(self, *, loc=None, ip=None) -> SymmetricBufferDevice: pointer_type = ir.Type.parse("!llvm.ptr") struct_type = _byval_struct_type(self.max_ranks) i64_type = ir.Type.parse("i64") - one = arith.constant( - value=ir.IntegerAttr.get(i64_type, 1), result=i64_type, loc=loc, ip=ip - ) - buffer = llvm.alloca( - res=pointer_type, - elem_type=struct_type, - array_size=one, - alignment=64, - loc=loc, - ip=ip, - ) + one = arith.constant(value=ir.IntegerAttr.get(i64_type, 1), result=i64_type, loc=loc, ip=ip) + buffer = llvm.alloca(res=pointer_type, elem_type=struct_type, array_size=one, alignment=64, loc=loc, ip=ip) for index, offset in enumerate(offsets): slot = llvm.getelementptr( - pointer_type, - buffer, - [], - [index], - i64_type, - no_wrap_flags="None", - loc=loc, - ip=ip, + pointer_type, buffer, [], [index], i64_type, no_wrap_flags="None", loc=loc, ip=ip ) llvm.store(self._as_int64(offset).ir_value(), slot, loc=loc, ip=ip) return SymmetricBufferDevice(buffer, self.max_ranks) @@ -148,12 +126,8 @@ def make_device_object(self, *, loc=None, ip=None) -> SymmetricBufferDevice: vector_type = ir.Type.parse(f"vector<{self.max_ranks}xi64>") vector = llvm.mlir_zero(vector_type, loc=loc, ip=ip) for index, offset in enumerate(offsets): - element_index = arith.constant( - value=ir.IntegerAttr.get(i32_type, index), result=i32_type, loc=loc, ip=ip - ) - vector = llvm.insertelement( - vector, self._as_int64(offset).ir_value(), element_index, loc=loc, ip=ip - ) + element_index = arith.constant(value=ir.IntegerAttr.get(i32_type, index), result=i32_type, loc=loc, ip=ip) + vector = llvm.insertelement(vector, self._as_int64(offset).ir_value(), element_index, loc=loc, ip=ip) return SymmetricBufferDevice(vector, self.max_ranks) @@ -163,14 +137,8 @@ def __init__(self, argument: SymmetricBufferHost) -> None: self._argument = argument offsets = tuple(argument.offsets) if len(offsets) != int(argument.max_ranks): - raise ValueError( - f"Expected {int(argument.max_ranks)} peer offsets, got {len(offsets)}." - ) - self._fields = ( - Int64(argument.base_address), - *(Int64(offset) for offset in offsets), - Int32(argument.rank), - ) + raise ValueError(f"Expected {int(argument.max_ranks)} peer offsets, got {len(offsets)}.") + self._fields = (Int64(argument.base_address), *(Int64(offset) for offset in offsets), Int32(argument.rank)) def __c_pointers__(self) -> list[Any]: pointers: list[Any] = [] @@ -195,9 +163,7 @@ def __new_from_mlir_values__(self, values: list[ir.Value]) -> SymmetricBufferHos rebuilt = [] for field in self._fields: field_value_count = len(get_mlir_types(field)) - rebuilt.append( - new_from_mlir_values(field, values[value_index : value_index + field_value_count]) - ) + rebuilt.append(new_from_mlir_values(field, values[value_index : value_index + field_value_count])) value_index += field_value_count if value_index != len(values): raise ValueError(f"Consumed {value_index} MLIR values, got {len(values)}.") diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm.py index 79fae23df062..910e753adb0f 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Metadata-push routing and fused non-deterministic token communication.""" import dataclasses @@ -8,6 +5,7 @@ from typing import Callable, ClassVar, Literal, Optional, Tuple, Union import cuda.bindings.driver as cuda + import cutlass import cutlass.cute as cute import cutlass.pipeline as pipeline @@ -19,8 +17,9 @@ from ...api import ImplDesc, KernelComponent, OptionalRequirement, ProblemDesc from ...helpers.device_workspace import DeviceWorkspace from ...helpers.dsl_helpers import mark_alignment, smem_exclusive_prefix -from ...helpers.flag_batch import make_flag_batch_tracker from ...helpers.iket_compat import iket +from ...helpers.flag_batch import make_flag_batch_tracker +from ...helpers.software_sync import NvlinkBarrier from ...helpers.ptx_helpers import ( cp_async_bulk_s2g, cp_reduce_async_bulk_add_bf16_s2g, @@ -33,25 +32,19 @@ tma_load_1d, ) from ...helpers.smem_workspace import SmemWorkspace -from ...helpers.software_sync import NvlinkBarrier from ...helpers.utils import ceil_div, round_up from ...quant_def import CombineFormat, QuantKind from ..token_protocol import TokenSrcMetadata from .symmetric_buffer import SymmetricBufferDevice + TokenBackMode = Literal["epi_warps", "standalone_warps", "reuse_dispatch_warps"] TokenBackScheduleMode = Literal["static", "atomic_counter"] @dataclasses.dataclass(frozen=True) class TokenCommArgs: - """Device views materialized inside the fused kernel region. - - ``pre_reduced_activation`` is the per-topk combine staging plane: TokenComm's own symmetric region, except - under in-kernel top-k reduction where it degenerates to a view of the caller's 2D output. That REDG - accumulates, so in that mode the incoming content is the caller's accumulation base -- zero, or a - shared-expert result. - """ + """Actual tensor and peer-mapping values for one communication invocation.""" activation: cute.Tensor activation_sf: cute.Tensor @@ -59,13 +52,25 @@ class TokenCommArgs: pre_reduced_activation_sf: Optional[cute.Tensor] peer_rank_ptr_mapper: SymmetricBufferDevice - def __extract_mlir_values__(self) -> list: - return [] - def __new_from_mlir_values__(self, values: list) -> "TokenCommArgs": - if values: - raise ValueError(f"TokenCommArgs expected no MLIR values, got {len(values)}.") - return self +@cute.jit +def _worker_coordinates( + warp_start: cutlass.Constexpr[int], + warp_count: cutlass.Constexpr[int], + warp_idx: Int32, + fused_warp_idx: Int32, + linear_cta_idx: Int32, + cta_ordinal: Optional[Int32], +): + """Return the local warp slot and dense worker index for a transfer operation.""" + if cutlass.const_expr(cta_ordinal is not None): + role_warp_idx = warp_idx - Int32(warp_start) + worker_cta = cta_ordinal + else: + # Legacy kernels select four transfer warps before calling the operation. + role_warp_idx = fused_warp_idx + worker_cta = linear_cta_idx + return role_warp_idx, worker_cta * Int32(warp_count) + role_warp_idx @dataclasses.dataclass(frozen=True) @@ -87,16 +92,12 @@ def from_packed(cls, packed: Union[Int64, Int32]) -> "_SortedElement": return cls(flat_topk_index=Int32(packed), topk_score=None) scratch = cute.make_rmem_tensor((2,), cutlass.Int32) cute.recast_tensor(scratch, cutlass.Int64)[0] = packed - return cls( - flat_topk_index=scratch[0], topk_score=cute.recast_tensor(scratch, cutlass.Float32)[1] - ) + return cls(flat_topk_index=scratch[0], topk_score=cute.recast_tensor(scratch, cutlass.Float32)[1]) @cute.jit def _copy_atom(dtype, num_bits_per_copy: int): - return cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), dtype, num_bits_per_copy=num_bits_per_copy - ) + return cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), dtype, num_bits_per_copy=num_bits_per_copy) class _MetadataPushRouter(KernelComponent): @@ -167,16 +168,13 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self._validate_router_configuration() self.expert_count_padded = round_up(self.expert_count, 4) self.expert_count_with_trash = self.expert_count_padded + 1 - self.router_elements_per_lane, self.router_data_cta_count = ( - self._router_launch_configuration() - ) + self.router_elements_per_lane, self.router_data_cta_count = self._router_launch_configuration() self.router_tokens_per_cta = self.router_elements_per_lane * self.router_warps_per_cta * 32 self.router_push_cta_count = ceil_div(self.expert_count, self.router_warps_per_cta) self.router_grid_cta_count = max(self.router_data_cta_count + 1, self.router_push_cta_count) if self.router_grid_cta_count > self.promised_launchable_sm_count: raise ValueError( - "Router grid exceeds promised_launchable_sm_count; " - "all metadata-push CTAs must be concurrently resident." + "Router grid exceeds promised_launchable_sm_count; all metadata-push CTAs must be concurrently resident." ) self.worst_case_token_count = self.worst_case_padded_tokens(self.token_padding_block) self._router_smem_workspace = self._build_router_smem_workspace() @@ -211,9 +209,7 @@ def _validate_router_configuration(self) -> None: if self.expert_count > 16384: raise NotImplementedError("TokenComm supports at most 16384 global experts.") if self.topk > self.expert_count: - raise ValueError( - f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}." - ) + raise ValueError(f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}.") @property def experts_per_rank(self) -> int: @@ -224,9 +220,7 @@ def worst_case_padded_tokens(self, block: int) -> int: routes_per_source_token = min(self.topk, self.experts_per_rank) route_capacity = source_token_capacity * routes_per_source_token active_expert_capacity = min(self.experts_per_rank, route_capacity) - route_budget_blocks = ( - active_expert_capacity + (route_capacity - active_expert_capacity) // block - ) + route_budget_blocks = active_expert_capacity + (route_capacity - active_expert_capacity) // block expert_bound_blocks = active_expert_capacity * int(ceil_div(source_token_capacity, block)) return min(route_budget_blocks, expert_bound_blocks) * block @@ -240,13 +234,8 @@ def next_power_of_two(value: int) -> int: maximum_data_cta_count = 128 maximum_supported_tokens = maximum_cta_capacity * maximum_data_cta_count if routed_token_capacity > maximum_supported_tokens: - raise NotImplementedError( - f"The router supports at most {maximum_supported_tokens} routed tokens per rank." - ) - cta_capacity = min( - maximum_cta_capacity, - next_power_of_two(max(routed_token_capacity, minimum_cta_capacity)), - ) + raise NotImplementedError(f"The router supports at most {maximum_supported_tokens} routed tokens per rank.") + cta_capacity = min(maximum_cta_capacity, next_power_of_two(max(routed_token_capacity, minimum_cta_capacity))) elements_per_lane = cta_capacity // (self.router_warps_per_cta * 32) data_cta_count = ceil_div(routed_token_capacity, cta_capacity) return elements_per_lane, data_cta_count @@ -256,18 +245,11 @@ def _build_router_smem_workspace(self) -> SmemWorkspace: workspace.register_mbarrier(self.router_helper_load_mbarrier_region, 1) overlay = workspace.create_overlay("nvlink.token_comm.router_smem.role") data_lifetime = overlay.add_lifetime("data_cta") + data_lifetime.register_tensor(self.router_data_histogram_region, cutlass.Int32, (self.expert_count_with_trash,)) data_lifetime.register_tensor( - self.router_data_histogram_region, cutlass.Int32, (self.expert_count_with_trash,) - ) - data_lifetime.register_tensor( - self.router_data_prefix_region, - cutlass.Int32, - (self.expert_count_with_trash,), - byte_alignment=16, - ) - data_lifetime.register_tensor( - self.router_data_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,) + self.router_data_prefix_region, cutlass.Int32, (self.expert_count_with_trash,), byte_alignment=16 ) + data_lifetime.register_tensor(self.router_data_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,)) data_lifetime.register_tensor( self.router_data_sorted_region, (cutlass.Int64 if self.apply_topk_at_fc1 else cutlass.Int32), @@ -276,10 +258,7 @@ def _build_router_smem_workspace(self) -> SmemWorkspace: ) if self.router_data_cta_count > 1: data_lifetime.register_tensor( - self.router_data_base_region, - cutlass.Int32, - (self.expert_count_padded,), - byte_alignment=16, + self.router_data_base_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 ) helper_lifetime = overlay.add_lifetime("helper_cta") @@ -291,22 +270,13 @@ def _build_router_smem_workspace(self) -> SmemWorkspace: byte_alignment=16, ) helper_lifetime.register_tensor( - self.router_helper_totals_region, - cutlass.Int32, - (self.expert_count_padded,), - byte_alignment=16, + self.router_helper_totals_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 ) helper_lifetime.register_tensor( - self.router_helper_prefix_region, - cutlass.Int32, - (self.expert_count_padded,), - byte_alignment=16, + self.router_helper_prefix_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 ) helper_lifetime.register_tensor( - self.router_helper_warp_totals_region, - cutlass.Int32, - (self.router_warps_per_cta,), - byte_alignment=16, + self.router_helper_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,), byte_alignment=16 ) workspace.finalize(max_bytes=self.router_smem_limit_bytes) return workspace @@ -329,34 +299,14 @@ def _register_router_workspace(self, workspace: DeviceWorkspace) -> None: stride=(self.expert_count_padded, 1), ) workspace.register( - self.sizes_region, - cutlass.Int32, - (self.expert_count_padded,), - buffer_space="shared", - reset="tail_reset", - ) - workspace.register( - self.sizes_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset" - ) - workspace.register( - self.metadata_ready_region, - cutlass.Int32, - (1,), - buffer_space="shared", - reset="tail_reset", - ) - workspace.register( - self.sorted_metadata_region, - cutlass.Int64, - (maximum_routed_tokens,), - buffer_space="local", + self.sizes_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="shared", reset="tail_reset" ) + workspace.register(self.sizes_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.metadata_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.sorted_metadata_region, cutlass.Int64, (maximum_routed_tokens,), buffer_space="local") if self.apply_topk_at_fc1: workspace.register( - self.sorted_scores_region, - cutlass.Float32, - (maximum_routed_tokens,), - buffer_space="local", + self.sorted_scores_region, cutlass.Float32, (maximum_routed_tokens,), buffer_space="local" ) workspace.register( self.token_src_metadata_region, @@ -367,43 +317,19 @@ def _register_router_workspace(self, workspace: DeviceWorkspace) -> None: ) if self.apply_topk_at_fc1: workspace.register( - self.fc1_topk_scores_region, - cutlass.Float32, - (self.worst_case_token_count,), - buffer_space="shared", + self.fc1_topk_scores_region, cutlass.Float32, (self.worst_case_token_count,), buffer_space="shared" ) + workspace.register(self.pool_expert_base_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local") workspace.register( - self.pool_expert_base_region, - cutlass.Int32, - (self.experts_per_rank,), - buffer_space="local", - ) - workspace.register( - self.source_expert_base_region, - cutlass.Int32, - (self.expert_count_padded,), - buffer_space="local", + self.source_expert_base_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="local" ) workspace.register( - self.push_destination_base_region, - cutlass.Int32, - (self.expert_count_padded,), - buffer_space="local", + self.push_destination_base_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="local" ) workspace.register( - self.sorted_metadata_ready_region, - cutlass.Int32, - (1,), - buffer_space="local", - reset="tail_reset", - ) - workspace.register( - self.push_table_ready_region, - cutlass.Int32, - (1,), - buffer_space="local", - reset="tail_reset", + self.sorted_metadata_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" ) + workspace.register(self.push_table_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset") if self.router_data_cta_count > 1: workspace.register( self.router_size_counter_region, @@ -413,18 +339,10 @@ def _register_router_workspace(self, workspace: DeviceWorkspace) -> None: reset="tail_reset", ) workspace.register( - self.router_histogram_done_region, - cutlass.Int32, - (1,), - buffer_space="local", - reset="tail_reset", + self.router_histogram_done_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" ) workspace.register( - self.source_base_ready_region, - cutlass.Int32, - (1,), - buffer_space="local", - reset="tail_reset", + self.source_base_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" ) def __extract_mlir_values__(self) -> list: @@ -524,16 +442,10 @@ def _router_single_cta( if self._router_linear_cta_idx < Int32(self.router_data_cta_count): block_thread_count = self.router_warps_per_cta * 32 trash_bucket = self.expert_count_padded - histogram = self._router_smem_workspace.tensor( - self.router_data_histogram_region, smem_base - ) + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) - warp_totals = self._router_smem_workspace.tensor( - self.router_data_warp_totals_region, smem_base - ) - sorted_elements = self._router_smem_workspace.tensor( - self.router_data_sorted_region, smem_base - ) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) zero_round_count = ceil_div(self.expert_count_with_trash, block_thread_count) for zero_round in cutlass.range_constexpr(zero_round_count): @@ -574,12 +486,7 @@ def _router_single_cta( iket.range_push("router.sort") self._sort_router_elements( - expert_registers, - within_expert_indices, - score_registers, - sorted_elements, - prefix, - topk_indices.dtype, + expert_registers, within_expert_indices, score_registers, sorted_elements, prefix, topk_indices.dtype ) cute.arch.sync_threads() iket.range_pop() @@ -588,10 +495,7 @@ def _router_single_cta( cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.sorted_metadata_ready_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" ) iket.range_pop() elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): @@ -605,16 +509,10 @@ def _router_multiple_ctas( block_thread_count = self.router_warps_per_cta * 32 trash_bucket = self.expert_count_padded expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) - histogram = self._router_smem_workspace.tensor( - self.router_data_histogram_region, smem_base - ) + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) - warp_totals = self._router_smem_workspace.tensor( - self.router_data_warp_totals_region, smem_base - ) - sorted_elements = self._router_smem_workspace.tensor( - self.router_data_sorted_region, smem_base - ) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) dump_base = self._router_smem_workspace.tensor(self.router_data_base_region, smem_base) zero_round_count = ceil_div(self.expert_count_with_trash, block_thread_count) @@ -635,17 +533,12 @@ def _router_multiple_ctas( expert = Int32(expert_round * block_thread_count) + self._router_thread_idx if expert < Int32(self.expert_count_padded): dump_base[expert] = Int32( - cute.arch.atomic_add( - size_counter + expert, histogram[expert], sem="relaxed", scope="gpu" - ) + cute.arch.atomic_add(size_counter + expert, histogram[expert], sem="relaxed", scope="gpu") ) cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.router_histogram_done_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.router_histogram_done_region), Int32(1), sem="release", scope="gpu" ) total_valid_routes = smem_exclusive_prefix( @@ -663,12 +556,7 @@ def _router_multiple_ctas( iket.range_pop() iket.range_push("router.sort") self._sort_router_elements( - expert_registers, - within_expert_indices, - score_registers, - sorted_elements, - prefix, - topk_indices.dtype, + expert_registers, within_expert_indices, score_registers, sorted_elements, prefix, topk_indices.dtype ) cute.arch.sync_threads() iket.range_pop() @@ -676,9 +564,7 @@ def _router_multiple_ctas( iket.range_push("router.wait_source_base") source_base_ready = self._device_workspace.ptr(self.source_base_ready_region) if self._router_thread_idx == Int32(0): - while cute.arch.load(source_base_ready, Int32, sem="acquire", scope="gpu") != Int32( - 1 - ): + while cute.arch.load(source_base_ready, Int32, sem="acquire", scope="gpu") != Int32(1): nanosleep(150) cute.arch.sync_threads() iket.range_pop() @@ -694,10 +580,7 @@ def _router_multiple_ctas( cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.sorted_metadata_ready_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" ) iket.range_pop() elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): @@ -714,9 +597,7 @@ def _router_helper_multiple_ctas(self, smem_base: cute.Pointer) -> None: block_thread_count = self.router_warps_per_cta * 32 totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) - warp_totals = self._router_smem_workspace.tensor( - self.router_helper_warp_totals_region, smem_base - ) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) size_counter = self._device_workspace.tensor(self.router_size_counter_region) source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) @@ -755,10 +636,7 @@ def _router_helper_multiple_ctas(self, smem_base: cute.Pointer) -> None: cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.source_base_ready_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.source_base_ready_region), Int32(1), sem="release", scope="gpu" ) publish_sizes() iket.range_pop() @@ -781,9 +659,7 @@ def _router_push_metadata(self) -> None: nanosleep(150) cute.arch.sync_threads() - global_expert = ( - self._router_linear_cta_idx * Int32(self.router_warps_per_cta) + self._router_warp_idx - ) + global_expert = self._router_linear_cta_idx * Int32(self.router_warps_per_cta) + self._router_warp_idx if global_expert < Int32(self.expert_count): sizes_by_rank = self._device_workspace.tensor(self.sizes_by_rank_region) source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) @@ -810,16 +686,12 @@ def _router_push_metadata(self) -> None: destination_position = destination_begin + route metadata = cute.arch.load(source_metadata + source_position, cutlass.Int64) stg_b64( - destination_metadata_address - + Int64(destination_position) * Int64(TokenSrcMetadata.nbytes), + destination_metadata_address + Int64(destination_position) * Int64(TokenSrcMetadata.nbytes), metadata, ) if cutlass.const_expr(self.apply_topk_at_fc1): score = cute.arch.load(source_scores + source_position, cutlass.Float32) - stg_f32( - destination_scores_address + Int64(destination_position) * Int64(4), - score, - ) + stg_f32(destination_scores_address + Int64(destination_position) * Int64(4), score) cute.arch.sync_threads() if self._router_thread_idx == Int32(0): @@ -832,10 +704,7 @@ def _router_push_metadata(self) -> None: destination_rank = Int32(rank_round * block_thread_count) + self._router_thread_idx if destination_rank < Int32(self.world_size): red_add_relaxed_sys_s32( - self._peer_rank_ptr_mapper.map( - metadata_ready_address, destination_rank, Int64(0) - ), - Int32(1), + self._peer_rank_ptr_mapper.map(metadata_ready_address, destination_rank, Int64(0)), Int32(1) ) @cute.jit @@ -852,8 +721,7 @@ def _broadcast_sizes_to_peers(self, smem_expert_counts: cute.Tensor) -> Callable destination_row_address = ( matrix_address + peer_offset - + Int64(Int32(self._router_local_rank) * Int32(self.expert_count_padded)) - * Int64(4) + + Int64(Int32(self._router_local_rank) * Int32(self.expert_count_padded)) * Int64(4) ) destination_row = cute.make_ptr( cutlass.Int32, destination_row_address, AddressSpace.gmem, assumed_align=16 @@ -863,9 +731,7 @@ def _broadcast_sizes_to_peers(self, smem_expert_counts: cute.Tensor) -> Callable ) with cute.arch.elect_one(): cp_async_bulk_s2g(destination_row, smem_expert_counts.iterator, row_bytes) - cp_reduce_async_bulk_add_u32_s2g( - destination_total, smem_expert_counts.iterator, row_bytes - ) + cp_reduce_async_bulk_add_u32_s2g(destination_total, smem_expert_counts.iterator, row_bytes) cute.arch.cp_async_bulk_commit_group() def finalize() -> None: @@ -880,8 +746,7 @@ def finalize() -> None: destination_rank = Int32(ready_round * block_thread_count) + self._router_thread_idx if destination_rank < Int32(self.world_size): red_add_relaxed_sys_s32( - self._peer_rank_ptr_mapper.map(ready_address, destination_rank, Int64(0)), - Int32(1), + self._peer_rank_ptr_mapper.map(ready_address, destination_rank, Int64(0)), Int32(1) ) return finalize @@ -892,19 +757,11 @@ def _compute_push_tables(self, smem_base: cute.Pointer) -> None: owner_expert_begin = Int32(self._router_local_rank) * Int32(self.experts_per_rank) matrix_bytes = self.world_size * self.expert_count_padded * 4 - size_matrix = self._router_smem_workspace.tensor( - self.router_helper_size_matrix_region, smem_base - ) - padded_totals = self._router_smem_workspace.tensor( - self.router_helper_totals_region, smem_base - ) + size_matrix = self._router_smem_workspace.tensor(self.router_helper_size_matrix_region, smem_base) + padded_totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) - warp_totals = self._router_smem_workspace.tensor( - self.router_helper_warp_totals_region, smem_base - ) - load_mbarrier = self._router_smem_workspace.ptr( - self.router_helper_load_mbarrier_region, smem_base - ) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) + load_mbarrier = self._router_smem_workspace.ptr(self.router_helper_load_mbarrier_region, smem_base) sizes = self._device_workspace.tensor(self.sizes_region) if self._router_thread_idx == Int32(0): @@ -913,9 +770,7 @@ def _compute_push_tables(self, smem_base: cute.Pointer) -> None: sizes_ready = self._device_workspace.ptr(self.sizes_ready_region) iket.range_push("router.wait_sizes_ready") if self._router_thread_idx == Int32(0): - while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="sys") != Int32( - self.world_size - ): + while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="sys") != Int32(self.world_size): nanosleep(150) cute.arch.mbarrier_init_fence() cute.arch.sync_threads() @@ -937,8 +792,7 @@ def _compute_push_tables(self, smem_base: cute.Pointer) -> None: if expert < Int32(self.expert_count_padded): expert_size = sizes[expert] padded_totals[expert] = ( - (expert_size + Int32(self.token_padding_block - 1)) - // Int32(self.token_padding_block) + (expert_size + Int32(self.token_padding_block - 1)) // Int32(self.token_padding_block) ) * Int32(self.token_padding_block) cute.arch.sync_threads() smem_exclusive_prefix( @@ -955,9 +809,7 @@ def _compute_push_tables(self, smem_base: cute.Pointer) -> None: for expert_round in cutlass.range_constexpr(local_expert_rounds): local_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx if local_expert < Int32(self.experts_per_rank): - pool_expert_base[local_expert] = ( - prefix[owner_expert_begin + local_expert] - prefix[owner_expert_begin] - ) + pool_expert_base[local_expert] = prefix[owner_expert_begin + local_expert] - prefix[owner_expert_begin] cute.arch.mbarrier_wait(load_mbarrier, 0) iket.range_pop() @@ -978,17 +830,12 @@ def _compute_push_tables(self, smem_base: cute.Pointer) -> None: for ring_position in cutlass.range_constexpr(self.world_size): source_rank = (destination_rank + Int32(ring_position)) % Int32(self.world_size) if Int32(ring_position) < local_ring_position: - source_ring_offset = ( - source_ring_offset + size_matrix[source_rank, global_expert] - ) + source_ring_offset = source_ring_offset + size_matrix[source_rank, global_expert] push_destination_base[global_expert] = destination_pool_base + source_ring_offset cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.push_table_ready_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.push_table_ready_region), Int32(1), sem="release", scope="gpu" ) iket.range_pop() @@ -1006,62 +853,46 @@ def _load_router_inputs( elements_per_thread = load_round_count * elements_per_vector topk_flat = cute.make_tensor(topk_indices.iterator, cute.make_layout((maximum_elements,))) - topk_vectors = cute.logical_divide( - cute.zipped_divide(topk_flat, (tile_span,)), (elements_per_vector, None) - ) + topk_vectors = cute.logical_divide(cute.zipped_divide(topk_flat, (tile_span,)), (elements_per_vector, None)) load_atom = _copy_atom(topk_indices.dtype, 128) expert_registers = cute.make_rmem_tensor((elements_per_thread,), cutlass.Int32) if cutlass.const_expr(topk_indices.dtype.width == 64): raw_indices = cute.make_rmem_tensor((elements_per_thread,), topk_indices.dtype) raw_vectors = cute.zipped_divide(raw_indices, (elements_per_vector,)) for load_round in cutlass.range_constexpr(load_round_count): - tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32( - elements_per_vector - ) + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) if tile_begin < actual_elements: cute.copy( load_atom, - mark_alignment( - topk_vectors[(None, self._router_grid_thread_idx), load_round], 16 - ), + mark_alignment(topk_vectors[(None, self._router_grid_thread_idx), load_round], 16), raw_vectors[None, load_round], ) else: expert_vectors = cute.zipped_divide(expert_registers, (elements_per_vector,)) for load_round in cutlass.range_constexpr(load_round_count): - tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32( - elements_per_vector - ) + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) if tile_begin < actual_elements: cute.copy( load_atom, - mark_alignment( - topk_vectors[(None, self._router_grid_thread_idx), load_round], 16 - ), + mark_alignment(topk_vectors[(None, self._router_grid_thread_idx), load_round], 16), expert_vectors[None, load_round], ) score_registers = None if cutlass.const_expr(self.apply_topk_at_fc1): score_registers = cute.make_rmem_tensor((elements_per_thread,), cutlass.Float32) - scores_flat = cute.make_tensor( - topk_scores.iterator, cute.make_layout((maximum_elements,)) - ) + scores_flat = cute.make_tensor(topk_scores.iterator, cute.make_layout((maximum_elements,))) score_vectors = cute.logical_divide( cute.zipped_divide(scores_flat, (tile_span,)), (elements_per_vector, None) ) score_atom = _copy_atom(cutlass.Float32, elements_per_vector * 32) score_register_vectors = cute.zipped_divide(score_registers, (elements_per_vector,)) for load_round in cutlass.range_constexpr(load_round_count): - tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32( - elements_per_vector - ) + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) if tile_begin < actual_elements: cute.copy( score_atom, - mark_alignment( - score_vectors[(None, self._router_grid_thread_idx), load_round], 16 - ), + mark_alignment(score_vectors[(None, self._router_grid_thread_idx), load_round], 16), score_register_vectors[None, load_round], ) @@ -1072,26 +903,21 @@ def _load_router_inputs( if cutlass.const_expr(topk_indices.dtype.width == 64): expert_registers[register_idx] = raw_indices_i32[2 * register_idx] token_idx, _ = self._router_value_coordinate(register_idx, topk_indices.dtype) - is_invalid = ( - expert_registers_u32[register_idx] >= cutlass.Uint32(self.expert_count) - ) | (token_idx >= actual_token_count) + is_invalid = (expert_registers_u32[register_idx] >= cutlass.Uint32(self.expert_count)) | ( + token_idx >= actual_token_count + ) if is_invalid: expert_registers[register_idx] = Int32(self.expert_count_padded) return expert_registers, score_registers @cute.jit - def _build_histogram( - self, expert_registers: cute.Tensor, histogram: cute.Tensor - ) -> cute.Tensor: + def _build_histogram(self, expert_registers: cute.Tensor, histogram: cute.Tensor) -> cute.Tensor: register_count = cute.size(expert_registers) within_expert_indices = cute.make_rmem_tensor((register_count,), cutlass.Int32) for register_idx in cutlass.range_constexpr(register_count): within_expert_indices[register_idx] = Int32( cute.arch.atomic_add( - histogram.iterator + expert_registers[register_idx], - Int32(1), - sem="relaxed", - scope="cta", + histogram.iterator + expert_registers[register_idx], Int32(1), sem="relaxed", scope="cta" ) ) cute.arch.sync_threads() @@ -1111,41 +937,28 @@ def _sort_router_elements( for register_idx in cutlass.range_constexpr(register_count): token_idx, topk_slot = self._router_value_coordinate(register_idx, topk_index_type) flat_topk_index = token_idx * Int32(self.topk) + topk_slot - destination = ( - expert_run_starts[expert_registers[register_idx]] - + within_expert_indices[register_idx] - ) + destination = expert_run_starts[expert_registers[register_idx]] + within_expert_indices[register_idx] if cutlass.const_expr(self.apply_topk_at_fc1): - sorted_elements[destination] = _SortedElement( - flat_topk_index, score_registers[register_idx] - ).pack() + sorted_elements[destination] = _SortedElement(flat_topk_index, score_registers[register_idx]).pack() else: sorted_elements[destination] = _SortedElement(flat_topk_index, None).pack() @cute.jit - def _router_value_coordinate( - self, register_idx: int, topk_index_type: type - ) -> Tuple[Int32, Int32]: + def _router_value_coordinate(self, register_idx: int, topk_index_type: type) -> Tuple[Int32, Int32]: elements_per_vector = 128 // topk_index_type.width - tile_span = ( - elements_per_vector * self.router_data_cta_count * self.router_warps_per_cta * 32 - ) + tile_span = elements_per_vector * self.router_data_cta_count * self.router_warps_per_cta * 32 flat_index = Int32( register_idx // elements_per_vector * tile_span + register_idx % elements_per_vector ) + self._router_grid_thread_idx * Int32(elements_per_vector) return (flat_index // Int32(self.topk), flat_index % Int32(self.topk)) @cute.jit - def _dump_contiguous_router_output( - self, sorted_elements: cute.Tensor, total_valid_routes: Int32 - ) -> None: + def _dump_contiguous_router_output(self, sorted_elements: cute.Tensor, total_valid_routes: Int32) -> None: block_thread_count = self.router_warps_per_cta * 32 metadata_address = self._device_workspace.ptr(self.sorted_metadata_region).toint() if cutlass.const_expr(self.apply_topk_at_fc1): score_address = self._device_workspace.ptr(self.sorted_scores_region).toint() - dump_round_count = (total_valid_routes + Int32(block_thread_count - 1)) // Int32( - block_thread_count - ) + dump_round_count = (total_valid_routes + Int32(block_thread_count - 1)) // Int32(block_thread_count) for dump_round in cutlass.range(dump_round_count, unroll=4): position = Int32(dump_round * block_thread_count) + self._router_thread_idx predicate = Int32(position < total_valid_routes) @@ -1155,11 +968,7 @@ def _dump_contiguous_router_output( src_token=(element.flat_topk_index // Int32(self.topk)), src_topk=(element.flat_topk_index % Int32(self.topk)), ) - stg_b64( - metadata_address + Int64(position) * Int64(TokenSrcMetadata.nbytes), - metadata.pack(), - predicate, - ) + stg_b64(metadata_address + Int64(position) * Int64(TokenSrcMetadata.nbytes), metadata.pack(), predicate) if cutlass.const_expr(self.apply_topk_at_fc1): stg_f32(score_address + Int64(position) * Int64(4), element.topk_score, predicate) @@ -1185,9 +994,7 @@ def _dump_router_output_by_expert( for route_round in cutlass.range(route_round_count, unroll=1): route = Int32(route_round) * Int32(32) + self._router_lane_idx predicate = Int32(route < run_length) - element = _SortedElement.from_packed( - sorted_elements[predicate * (run_begin + route)] - ) + element = _SortedElement.from_packed(sorted_elements[predicate * (run_begin + route)]) output_position = dump_begin + route metadata = TokenSrcMetadata( src_rank=Int32(self._router_local_rank), @@ -1200,22 +1007,14 @@ def _dump_router_output_by_expert( predicate, ) if cutlass.const_expr(self.apply_topk_at_fc1): - stg_f32( - score_address + Int64(output_position) * Int64(4), - element.topk_score, - predicate, - ) + stg_f32(score_address + Int64(output_position) * Int64(4), element.topk_score, predicate) @cute.jit - def local_expert_sizes( - self, device_workspace: DeviceWorkspace, local_rank: Int32 - ) -> cute.Tensor: + def local_expert_sizes(self, device_workspace: DeviceWorkspace, local_rank: Int32) -> cute.Tensor: """Return this rank's contiguous expert-size view.""" sizes = device_workspace.tensor(self.sizes_region) expert_begin = local_rank * Int32(self.experts_per_rank) - return cute.make_tensor( - sizes.iterator + expert_begin, cute.make_layout((self.experts_per_rank,)) - ) + return cute.make_tensor(sizes.iterator + expert_begin, cute.make_layout((self.experts_per_rank,))) @property def metadata_ready_target(self) -> int: @@ -1234,23 +1033,17 @@ def token_src_metadata_pointer(self, device_workspace: DeviceWorkspace) -> cute. return device_workspace.ptr(self.token_src_metadata_region) @cute.jit - def wait_for_sizes_ready( - self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000 - ) -> None: + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: thread_idx, _, _ = cute.arch.thread_idx() lane_idx = thread_idx % Int32(32) if lane_idx == Int32(0): sizes_ready = device_workspace.ptr(self.sizes_ready_region) - while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="sys") != Int32( - self.world_size - ): + while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="sys") != Int32(self.world_size): nanosleep(sleep_cycles) cute.arch.sync_warp() @cute.jit - def wait_for_metadata_ready( - self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000 - ) -> None: + def wait_for_metadata_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: thread_idx, _, _ = cute.arch.thread_idx() lane_idx = thread_idx % Int32(32) if lane_idx == Int32(0): @@ -1301,6 +1094,7 @@ class TokenCommNonDeterministic(KernelComponent): expert_sizes_smem_region = "nvlink.token_comm.main_smem.expert_sizes" token_in_activation_smem_region = "nvlink.token_comm.main_smem.token_in_activation" # nosec B105 token_in_sf_smem_region = "nvlink.token_comm.main_smem.token_in_sf" # nosec B105 + token_in_sf_group_region = "nvlink.token_comm.main_smem.token_in_sf_group" # nosec B105 token_back_activation_smem_region = "nvlink.token_comm.main_smem.token_back_activation" # nosec B105 token_back_sf_smem_region = "nvlink.token_comm.main_smem.token_back_sf" # nosec B105 @@ -1328,8 +1122,15 @@ def impl_desc_require(cls) -> dict[str, type]: "token_in_flag_batch": int, "token_back_mode": str, "token_back_schedule_mode": str, + "token_back_ready_granularity": OptionalRequirement(str), "reduce_topk_in_kernel": bool, "router_smem_limit_bytes": OptionalRequirement(int), + "comm_cta_count": OptionalRequirement(int), + "comm_warps": OptionalRequirement(tuple), + "token_back_on_compute_ctas": OptionalRequirement(bool), + "comm_stage_depths": OptionalRequirement(tuple), + "token_back_chunk_bytes": OptionalRequirement(int), + "dispatch_sf_group": OptionalRequirement(int), } def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: @@ -1352,20 +1153,80 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.token_in_flag_batch = impl_desc["token_in_flag_batch"] self.token_back_mode: TokenBackMode = impl_desc["token_back_mode"] self.token_back_schedule_mode: TokenBackScheduleMode = impl_desc["token_back_schedule_mode"] + self.token_back_ready_granularity = impl_desc.get("token_back_ready_granularity", "expert") self.reduce_topk_in_kernel = impl_desc["reduce_topk_in_kernel"] self.router_smem_limit_bytes = impl_desc.get("router_smem_limit_bytes", 227 * 1024) + comm_cta_count = impl_desc.get("comm_cta_count", 0) + if ( + isinstance(comm_cta_count, bool) + or not isinstance(comm_cta_count, int) + or not 0 <= comm_cta_count < self.promised_launchable_sm_count + ): + raise ValueError("comm_cta_count must be a nonnegative integer leaving at least one compute CTA.") + self.cta_specialized = comm_cta_count > 0 + if self.cta_specialized: + comm_warps = impl_desc.get("comm_warps") + if ( + not isinstance(comm_warps, tuple) + or len(comm_warps) != 2 + or any(isinstance(w, bool) or not isinstance(w, int) or w <= 0 for w in comm_warps) + ): + raise ValueError("Dedicated communication CTAs require comm_warps as two positive integers.") + self.dispatch_warp_count, self.combine_warp_count = comm_warps + back_on_compute = impl_desc.get("token_back_on_compute_ctas", False) + self.dispatch_warp_start = 0 + self.combine_warp_start = 8 if back_on_compute else self.dispatch_warp_count + self.dispatch_cta_count = comm_cta_count + self.combine_cta_count = ( + self.promised_launchable_sm_count - comm_cta_count if back_on_compute else comm_cta_count + ) + else: + # FC12 occupies warps 0..7. Legacy primitive callers retain their + # modulo-four numbering by omitting the explicit CTA ordinal. + self.dispatch_warp_start = 8 + self.combine_warp_start = 12 if self.token_back_mode == "standalone_warps" else 8 # nosec B105 + self.dispatch_warp_count = self.transfer_warp_count + self.combine_warp_count = self.transfer_warp_count + self.dispatch_cta_count = self.promised_launchable_sm_count + self.combine_cta_count = self.promised_launchable_sm_count + self.comm_stage_depths = tuple(impl_desc.get("comm_stage_depths", (1, 1))) + self.token_back_chunk_bytes = impl_desc.get("token_back_chunk_bytes", self.standalone_chunk_bytes) + # Tokens a token_in warp batches before writing their scale rows; see _token_in_multi_stage. + self.dispatch_sf_group = impl_desc.get("dispatch_sf_group", 1) self._validate_configuration() self._router = _MetadataPushRouter(problem_desc, impl_desc) - self._nvlink_barrier = NvlinkBarrier( - world_size=self.world_size, barrier_id=self.grid_sync_barrier_id - ) + self._nvlink_barrier = NvlinkBarrier(world_size=self.world_size, barrier_id=self.grid_sync_barrier_id) self._device_workspace = None self._token_comm_args = None self._local_rank = None self._linear_cta_idx = None self._transfer_warp_idx = None self._lane_idx = None + self._dispatch_warp_idx = None + self._combine_warp_idx = None + self._dispatch_worker_idx = None + self._combine_worker_idx = None + + @property + def dispatch_stage_depth(self) -> int: + """Tokens one token_in warp keeps in flight; ``1`` is the single-buffered loop.""" + return self.comm_stage_depths[0] + + @property + def combine_stage_depth(self) -> int: + """Output chunks one token_back warp keeps in flight.""" + return self.comm_stage_depths[1] + + @property + def combine_output_token_bytes(self) -> int: + """Bytes one combine output row occupies on the wire.""" + return self.hidden_size * int(self.combine_format.act_dtype.width) // 8 + + @property + def sf_k_group_count(self) -> int: + """Scale vectors of one token, in groups of four -- the width the destination interleaves at.""" + return self.activation_sf_hidden_padded // 4 def _validate_configuration(self) -> None: positive_fields = ( @@ -1391,35 +1252,40 @@ def _validate_configuration(self) -> None: if self.expert_count > 16384: raise NotImplementedError("TokenComm supports at most 16384 global experts.") if self.topk > self.expert_count: - raise ValueError( - f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}." - ) + raise ValueError(f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}.") if self.token_back_mode not in ("epi_warps", "standalone_warps", "reuse_dispatch_warps"): raise ValueError(f"Unsupported token_back_mode {self.token_back_mode!r}.") if self.token_back_schedule_mode not in ("static", "atomic_counter"): raise ValueError( f"token_back_schedule_mode must be static or atomic_counter, got {self.token_back_schedule_mode!r}." ) + if self.token_back_ready_granularity not in ("expert", "token_tile"): + raise ValueError(f"Unsupported token_back_ready_granularity {self.token_back_ready_granularity!r}.") + if self.token_back_ready_granularity == "token_tile" and ( # nosec B105 + self.token_back_mode not in ("standalone_warps", "reuse_dispatch_warps") + or self.token_back_schedule_mode != "atomic_counter" # nosec B105 + ): + raise ValueError("token_tile return-ready requires standalone_warps or reuse_dispatch_warps with atomic_counter.") if not 1 <= self.token_in_flag_batch <= 32: - raise ValueError( - f"token_in_flag_batch must be in [1, 32], got {self.token_in_flag_batch}." - ) + raise ValueError(f"token_in_flag_batch must be in [1, 32], got {self.token_in_flag_batch}.") if self.tokens_per_fc1_ready_slot % self.token_padding_block != 0: raise ValueError("tokens_per_fc1_ready_slot must be divisible by token_padding_block.") if self.token_back_enabled and self.fc2_done_signals_per_token_tile <= 0: - raise ValueError( - "fc2_done_signals_per_token_tile must be positive when token-back is enabled." - ) + raise ValueError("fc2_done_signals_per_token_tile must be positive when token-back is enabled.") if self.reduce_topk_in_kernel and self.combine_format.act_dtype is not cutlass.BFloat16: raise ValueError("In-kernel top-k reduction requires BF16 combine data.") element_block = self.activation_sf_vector_size * 4 if self.hidden_size % element_block != 0: - raise ValueError( - f"{self.quant_kind} requires hidden_size divisible by {element_block}." - ) + raise ValueError(f"{self.quant_kind} requires hidden_size divisible by {element_block}.") if self.sf_padding_block % 128 != 0: raise ValueError("sf_padding_block must be a multiple of 128.") + if self.token_back_enabled and self.combine_stage_depth > 1: + if self.token_back_schedule_mode != "static": # nosec B105 + raise ValueError("Multi-stage token-back requires static scheduling.") + if self.combine_format.is_quantized: + raise ValueError("Multi-stage token-back supports BF16 combine only.") + @property def experts_per_rank(self) -> int: return self.expert_count // self.world_size @@ -1476,10 +1342,7 @@ def worst_case_sf_token_count(self) -> int: @property def max_fc1_ready_slot_count(self) -> int: - return ( - self._router.worst_case_padded_tokens(self.tokens_per_fc1_ready_slot) - // self.tokens_per_fc1_ready_slot - ) + return self._router.worst_case_padded_tokens(self.tokens_per_fc1_ready_slot) // self.tokens_per_fc1_ready_slot @property def router_smem_workspace(self) -> SmemWorkspace: @@ -1550,15 +1413,11 @@ def launch_router( ) @cute.jit - def local_expert_sizes( - self, device_workspace: DeviceWorkspace, local_rank: Int32 - ) -> cute.Tensor: + def local_expert_sizes(self, device_workspace: DeviceWorkspace, local_rank: Int32) -> cute.Tensor: return self._router.local_expert_sizes(device_workspace, local_rank) @cute.jit - def wait_for_sizes_ready( - self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000 - ) -> None: + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: self._router.wait_for_sizes_ready(device_workspace, sleep_cycles) @cute.jit @@ -1577,18 +1436,40 @@ def assign_device_members( token_comm_args: TokenCommArgs, local_rank: Int32, linear_cta_idx: Int32, + dispatch_cta_ordinal: Optional[Int32] = None, + combine_cta_ordinal: Optional[Int32] = None, ) -> None: + if cutlass.const_expr( + self.cta_specialized and (dispatch_cta_ordinal is None or combine_cta_ordinal is None) + ): + raise ValueError("Specialized token communication requires explicit dispatch and combine CTA ordinals.") self._device_workspace = device_workspace self._token_comm_args = token_comm_args self._local_rank = local_rank self._linear_cta_idx = linear_cta_idx thread_idx, _, _ = cute.arch.thread_idx() - transfer_thread_idx = thread_idx % Int32(self.transfer_thread_count) - self._transfer_warp_idx = cute.arch.make_warp_uniform(transfer_thread_idx // Int32(32)) - self._lane_idx = transfer_thread_idx % Int32(32) - self._nvlink_barrier.assign_device_members( - device_workspace, token_comm_args.peer_rank_ptr_mapper - ) + warp_idx = cute.arch.make_warp_uniform(thread_idx // Int32(32)) + self._transfer_warp_idx = cute.arch.make_warp_uniform( + (thread_idx % Int32(self.transfer_thread_count)) // Int32(32) + ) + self._lane_idx = thread_idx % Int32(32) + self._dispatch_warp_idx, self._dispatch_worker_idx = _worker_coordinates( + self.dispatch_warp_start, + self.dispatch_warp_count, + warp_idx, + self._transfer_warp_idx, + linear_cta_idx, + dispatch_cta_ordinal, + ) + self._combine_warp_idx, self._combine_worker_idx = _worker_coordinates( + self.combine_warp_start, + self.combine_warp_count, + warp_idx, + self._transfer_warp_idx, + linear_cta_idx, + combine_cta_ordinal, + ) + self._nvlink_barrier.assign_device_members(device_workspace, token_comm_args.peer_rank_ptr_mapper) def remove_device_members(self) -> None: self._nvlink_barrier.remove_device_members() @@ -1598,6 +1479,10 @@ def remove_device_members(self) -> None: self._linear_cta_idx = None self._transfer_warp_idx = None self._lane_idx = None + self._dispatch_warp_idx = None + self._combine_warp_idx = None + self._dispatch_worker_idx = None + self._combine_worker_idx = None def __extract_mlir_values__(self) -> list: return [] @@ -1643,9 +1528,7 @@ def _register_main_workspace(self, workspace: DeviceWorkspace) -> None: buffer_space="local", byte_alignment=128, ) - activation_sf_element_count = ( - self.worst_case_sf_token_count * self.activation_sf_hidden_padded - ) + activation_sf_element_count = self.worst_case_sf_token_count * self.activation_sf_hidden_padded workspace.register( self.fc1_activation_sf_region, self.activation_sf_dtype, @@ -1655,12 +1538,13 @@ def _register_main_workspace(self, workspace: DeviceWorkspace) -> None: ) if self.token_back_enabled: + fc2_done_slots = ( + self.max_fc1_ready_slot_count + if self.token_back_ready_granularity == "token_tile" # nosec B105 + else self.experts_per_rank + ) workspace.register( - self.fc2_done_region, - cutlass.Int32, - (self.experts_per_rank,), - buffer_space="local", - reset="tail_reset", + self.fc2_done_region, cutlass.Int32, (fc2_done_slots,), buffer_space="local", reset="tail_reset" ) if self.token_back_push_data: fc2_element_count = self.worst_case_token_count * self.hidden_size @@ -1682,11 +1566,7 @@ def _register_main_workspace(self, workspace: DeviceWorkspace) -> None: ) if self.token_back_enabled and self.token_back_schedule_mode == "atomic_counter": # nosec B105 workspace.register( - self.token_back_schedule_region, - cutlass.Int32, - (1,), - buffer_space="local", - reset="tail_reset", + self.token_back_schedule_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" ) if not self.reduce_topk_in_kernel: @@ -1709,9 +1589,16 @@ def _register_main_workspace(self, workspace: DeviceWorkspace) -> None: ) def register_smem_regions(self, workspace: SmemWorkspace) -> None: - workspace.register_mbarrier(self.token_in_mbarrier_region, self.transfer_warp_count) + """Legacy colocated placement, preserving its declaration order and reuse overlay.""" + workspace.register_mbarrier( + self.token_in_mbarrier_region, + self.dispatch_warp_count * self.dispatch_stage_depth, + ) if self.token_back_enabled: - workspace.register_mbarrier(self.token_back_mbarrier_region, self.transfer_warp_count) + workspace.register_mbarrier( + self.token_back_mbarrier_region, + self.combine_warp_count * self.combine_stage_depth, + ) workspace.register_tensor( self.expert_sizes_smem_region, cutlass.Int32, @@ -1719,76 +1606,101 @@ def register_smem_regions(self, workspace: SmemWorkspace) -> None: byte_alignment=16, ) transfer_overlay = workspace.create_overlay("nvlink.token_comm.main_smem.transfer") - token_in_lifetime = transfer_overlay.add_lifetime("token_in") - token_in_lifetime.register_tensor( + self._register_token_in_staging(transfer_overlay.add_lifetime("token_in")) + if not self.token_back_enabled: + return + if self.token_back_mode == "standalone_warps": # nosec B105 + back_overlay = workspace.create_overlay("nvlink.token_comm.main_smem.standalone_token_back") + else: + back_overlay = transfer_overlay + self._register_token_back_staging(back_overlay.add_lifetime("token_back")) + + def _register_token_in_staging(self, staging_lifetime) -> None: + """Declare input tensors in the caller-selected lifetime.""" + dispatch_rings = self.dispatch_warp_count + dispatch_depth = self.dispatch_stage_depth + sf_padded = self.activation_sf_hidden_padded + if dispatch_depth == 1: + activation_shape = (dispatch_rings, self.hidden_size) + sf_shape = (dispatch_rings, (self.activation_sf_vector_size, sf_padded)) + sf_stride = (sf_padded, (0, 1)) + else: + activation_shape = (dispatch_rings, dispatch_depth, self.hidden_size) + sf_shape = ( + dispatch_rings, + dispatch_depth, + (self.activation_sf_vector_size, sf_padded), + ) + sf_stride = (dispatch_depth * sf_padded, sf_padded, (0, 1)) + staging_lifetime.register_tensor( self.token_in_activation_smem_region, self.activation_dtype, - (self.transfer_warp_count, self.hidden_size), + activation_shape, byte_alignment=16, ) - token_in_lifetime.register_tensor( + staging_lifetime.register_tensor( self.token_in_sf_smem_region, self.activation_sf_dtype, - ( - self.transfer_warp_count, - (self.activation_sf_vector_size, self.activation_sf_hidden_padded), - ), - stride=(self.activation_sf_hidden_padded, (0, 1)), + sf_shape, + stride=sf_stride, byte_alignment=16, ) + if self.dispatch_stage_depth > 1: + # One k-group of a token group, laid out the way the destination wants it. Shares the token_in + # lifetime because it is only live between a group's first arrival and its scale store. The + # condition is the depth rather than the group size, because it is the depth that selects the + # grouped loop -- a group of one still stages its scales here. + staging_lifetime.register_tensor( + self.token_in_sf_group_region, + self.activation_sf_dtype, + (dispatch_rings, self.sf_k_group_count, 4 * self.dispatch_sf_group), + byte_alignment=16, + ) + + def _register_token_back_staging(self, staging_lifetime) -> None: + """Declare output tensors in the caller-selected lifetime.""" if not self.token_back_enabled: return - - if self.token_back_mode == "standalone_warps": # nosec B105 - token_back_lifetime = workspace.create_overlay( - "nvlink.token_comm.main_smem.standalone_token_back" - ).add_lifetime("token_back") - else: - token_back_lifetime = transfer_overlay.add_lifetime("token_back") - + combine_rings = self.combine_warp_count + combine_depth = self.combine_stage_depth + sf_padded = self.activation_sf_hidden_padded if self.token_back_mode == "standalone_warps": # nosec B105 - available_bytes_per_warp = self.standalone_chunk_bytes + available_bytes_per_warp = self.token_back_chunk_bytes else: - activation_bytes = self.bytes_per_token - sf_bytes = self.activation_sf_hidden_padded * int(self.activation_sf_dtype.width) // 8 - available_bytes_per_warp = activation_bytes + sf_bytes + available_bytes_per_warp = self.bytes_per_token + sf_padded * int(self.activation_sf_dtype.width) // 8 if self.token_back_push_data: - bytes_per_output_token = ( - self.hidden_size * int(self.combine_format.act_dtype.width) // 8 - ) + bytes_per_output_token = self.hidden_size * int(self.combine_format.act_dtype.width) // 8 if self.token_back_mode == "standalone_warps": # nosec B105 - chunk_bytes = self.standalone_chunk_bytes + chunk_bytes = self.token_back_chunk_bytes elif available_bytes_per_warp < bytes_per_output_token: chunk_bytes = self.bytes_per_token else: chunk_bytes = bytes_per_output_token - if ( - self.token_back_mode != "standalone_warps" # nosec B105 - and bytes_per_output_token % chunk_bytes != 0 - ): + if self.token_back_mode != "standalone_warps" and bytes_per_output_token % chunk_bytes != 0: # nosec B105 raise ValueError("Token-back data chunk bytes must divide one row.") chunk_elements = chunk_bytes * 8 // int(self.combine_format.act_dtype.width) - token_back_lifetime.register_tensor( + shape = ( + (combine_rings, chunk_elements) + if combine_depth == 1 + else (combine_rings, combine_depth, chunk_elements) + ) + staging_lifetime.register_tensor( self.token_back_activation_smem_region, self.combine_format.act_dtype, - (self.transfer_warp_count, chunk_elements), + shape, byte_alignment=16, ) if self.token_back_push_sf: - sf_row_bytes = ( - self.combine_sf_hidden_padded * int(self.combine_format.scale_dtype.width) // 8 - ) + sf_row_bytes = self.combine_sf_hidden_padded * int(self.combine_format.scale_dtype.width) // 8 if sf_row_bytes > available_bytes_per_warp: raise ValueError("Token-back scale row exceeds its per-warp stage.") - token_back_lifetime.register_tensor( + combine_sf = self.combine_sf_hidden_padded + staging_lifetime.register_tensor( self.token_back_sf_smem_region, self.combine_format.scale_dtype, - ( - self.transfer_warp_count, - (self.combine_format.scale_block, self.combine_sf_hidden_padded), - ), - stride=(self.combine_sf_hidden_padded, (0, 1)), + (combine_rings, (self.combine_format.scale_block, combine_sf)), + stride=(combine_sf, (0, 1)), byte_alignment=16, ) @@ -1800,9 +1712,7 @@ def fc1_ready_counter_pointer(self, device_workspace: DeviceWorkspace) -> cute.P def fc1_activation_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: return cute.make_tensor( device_workspace.ptr(self.fc1_activation_region), - cute.make_layout( - (self.worst_case_token_count, self.hidden_size), stride=(self.hidden_size, 1) - ), + cute.make_layout((self.worst_case_token_count, self.hidden_size), stride=(self.hidden_size, 1)), ) @cute.jit @@ -1810,9 +1720,7 @@ def fc1_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> cute.Te layout = tile_atom_to_shape_SF( (self.worst_case_sf_token_count, self.hidden_size, 1), self.activation_sf_vector_size ) - return cute.make_tensor( - device_workspace.ptr(self.fc1_activation_sf_region), cute.select(layout, mode=[0, 1]) - ) + return cute.make_tensor(device_workspace.ptr(self.fc1_activation_sf_region), cute.select(layout, mode=[0, 1])) @cute.jit def fc2_done_counter_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: @@ -1827,8 +1735,7 @@ def fc2_activation_tensor(self, device_workspace: DeviceWorkspace) -> Optional[c return cute.make_tensor( device_workspace.ptr(self.fc2_activation_region), cute.make_layout( - (self.worst_case_token_count, 1, self.hidden_size), - stride=(self.hidden_size, self.hidden_size, 1), + (self.worst_case_token_count, 1, self.hidden_size), stride=(self.hidden_size, self.hidden_size, 1) ), ) @@ -1842,40 +1749,375 @@ def fc2_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optiona ( self.worst_case_token_count, 1, - ( - self.combine_format.scale_block, - self.hidden_size // self.combine_format.scale_block, - ), + (self.combine_format.scale_block, self.hidden_size // self.combine_format.scale_block), ), stride=(self.combine_sf_hidden_padded, self.combine_sf_hidden_padded, (0, 1)), ), ) @cute.jit - def pre_reduced_activation_tensor( - self, device_workspace: DeviceWorkspace - ) -> Optional[cute.Tensor]: + def pre_reduced_activation_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: """The (tokens, topk, hidden) combine staging plane, or None under in-kernel top-k reduction.""" if cutlass.const_expr(self.reduce_topk_in_kernel): return None return device_workspace.tensor(self.pre_reduced_activation_region) @cute.jit - def pre_reduced_activation_sf_tensor( - self, device_workspace: DeviceWorkspace - ) -> Optional[cute.Tensor]: + def pre_reduced_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: """The scale plane parallel to ``pre_reduced_activation_tensor``; only quantized wire formats carry one.""" if cutlass.const_expr(self.reduce_topk_in_kernel or not self.combine_format.is_quantized): return None return device_workspace.tensor(self.pre_reduced_activation_sf_region) + @cute.jit + def _token_in_multi_stage( + self, + smem_workspace, + smem_base, + owned_sizes, + pool_expert_bases, + transfer_warp_idx, + lane_idx, + global_warp_idx, + global_warp_count, + ) -> None: + """Pull tokens through a depth-D ring, writing their scale rows a group at a time. + + Where a scale byte lands: ``r * 512 + (t % 32) * 16 + (t // 32) * 4 + k``. One 128-token atom block + spreads over 112 such k-groups, and inside each one the 512 bytes are laid out by token, not in + token order:: + + k-group r: r*512 r*512 + 511 + |<------------------ 512 B ------------------------>| + +---------+---------+---------+ +---------+ + | j = 0 | j = 1 | j = 2 | ... | j = 31 | 16 B per j + +---------+---------+---------+ +---------+ + | t | t | t | t | + 0 32 64 96 4 B per token + + So token ``t`` owns four bytes here, four more 512 B later, 112 times over: written per token that + is 112 scattered stores against the ONE bulk copy its 3584-byte activation takes. The cost is the + store count, not the bytes. + + The four tokens sharing a ``j`` are exactly one contiguous sixteen. Staging ``{j, j+32, j+64, j+96}`` + together and writing that run turns the 112 four-byte stores into 112 sixteen-byte ones, which is why + the worker space strides over ``(block, j)`` pairs rather than tokens. + + Those four stages stay pinned until the run is written, so the ring is consumed a group at a time and + the depth must be a multiple of the group. + """ + group_size = self.dispatch_sf_group + depth = self.dispatch_stage_depth + group_stage_sets = depth // group_size + k_groups = self.sf_k_group_count + # One 128-token atom block is 112 atoms of 512 bytes, and that is also the stride between blocks. + sf_block_bytes = k_groups * 512 + + token_in_activation = smem_workspace.tensor(self.token_in_activation_smem_region, smem_base) + token_in_sf = smem_workspace.tensor(self.token_in_sf_smem_region, smem_base) + sf_group = smem_workspace.tensor(self.token_in_sf_group_region, smem_base) + warp_mbarrier = smem_workspace.ptr(self.token_in_mbarrier_region, smem_base) + transfer_warp_idx * Int32(depth) + if lane_idx == Int32(0): + for init_stage in cutlass.range_constexpr(depth): + cute.arch.mbarrier_init(warp_mbarrier + Int32(init_stage), 1) + cute.arch.sync_warp() + + fc1_activation_pointer = self._device_workspace.ptr(self.fc1_activation_region) + fc1_activation_sf_pointer = self._device_workspace.ptr(self.fc1_activation_sf_region) + fc1_ready_counter = self._device_workspace.ptr(self.fc1_ready_region) + activation_bytes = self.bytes_per_token + scatter_atom = _copy_atom( + self.activation_sf_dtype, + 4 * int(self.activation_sf_dtype.width), + ) + store_atom = _copy_atom( + self.activation_sf_dtype, + 4 * group_size * int(self.activation_sf_dtype.width), + ) + + first_group_in_this_expert = global_warp_idx + expert_group_begin = Int32(0) + expert_sf_begin = Int32(0) + expert_ready_slot_begin = Int32(0) + pull_phase_mask = Int32(0) + flag_tracker = make_flag_batch_tracker( + use_async=self.token_in_flag_batch == 1, + flag_address=Int64(0), + accumulated_flags=Int32(0), + phase=Int32(0), + thread_idx=lane_idx, + ) + + local_expert = Int32(0) + while local_expert < Int32(self.experts_per_rank): + expert_token_count = owned_sizes[local_expert] + block_count = (expert_token_count + Int32(127)) // Int32(128) + expert_group_end = expert_group_begin + block_count * Int32(32) + group_count = Int32(0) + if first_group_in_this_expert < expert_group_end: + group_count = ( + expert_group_end - first_group_in_this_expert + global_warp_count - Int32(1) + ) // global_warp_count + pool_expert_base = pool_expert_bases[local_expert] + + # One group in flight before the first is consumed. The stage cursor restarts at 0 every expert + # because every group issued here is consumed here, so the ring holds nothing unconsumed at the + # boundary; the per-stage mbarrier phases do NOT restart, since a stage's parity is the count of + # waits it has served. + if group_count > Int32(0): + self._issue_token_group( + smem_workspace, + smem_base, + Int32(0), + first_group_in_this_expert, + expert_group_begin, + global_warp_count, + expert_token_count, + pool_expert_base, + Int32(0), + transfer_warp_idx, + lane_idx, + ) + + stage_base = Int32(0) + for group_round in cutlass.range(group_count, unroll=1): + next_stage_base = stage_base + Int32(group_size) + # Both arms assign: a value carried out of a staged conditional has to be written on + # every path, so the else is required rather than redundant. + if next_stage_base >= Int32(depth): + next_stage_base = next_stage_base - Int32(depth) + else: + next_stage_base = next_stage_base + if cutlass.const_expr(group_stage_sets > 1): + # Spare stages exist, so the next group's reads run under this group's drain. + if group_round + Int32(1) < group_count: + self._issue_token_group( + smem_workspace, + smem_base, + group_round + Int32(1), + first_group_in_this_expert, + expert_group_begin, + global_warp_count, + expert_token_count, + pool_expert_base, + next_stage_base, + transfer_warp_idx, + lane_idx, + ) + + group_in_this_expert = first_group_in_this_expert + group_round * global_warp_count - expert_group_begin + block_idx = group_in_this_expert // Int32(32) + column_in_block = group_in_this_expert - block_idx * Int32(32) + + for token_in_group in cutlass.range_constexpr(group_size): + token_in_expert = block_idx * Int32(128) + column_in_block + Int32(32 * token_in_group) + if token_in_expert < expert_token_count: + stage_idx = stage_base + Int32(token_in_group) + cute.arch.mbarrier_wait( + warp_mbarrier + stage_idx, + (pull_phase_mask >> stage_idx) & Int32(1), + ) + pull_phase_mask = pull_phase_mask ^ (Int32(1) << stage_idx) + destination_activation = cute.make_ptr( + self.activation_dtype, + fc1_activation_pointer.toint() + + Int64(pool_expert_base + token_in_expert) * Int64(self.bytes_per_token), + AddressSpace.gmem, + assumed_align=16, + ) + if lane_idx == Int32(0): + cp_async_bulk_s2g( + destination_activation, + token_in_activation[transfer_warp_idx, stage_idx, None].iterator, + Int32(activation_bytes), + ) + cute.arch.cp_async_bulk_commit_group() + cute.arch.sync_warp() + # Transpose this token's contiguous scale row into the group buffer's destination + # order. These are SMEM stores; the global ones happen once per group below. + stage_sf = token_in_sf[transfer_warp_idx, stage_idx, (None, None)] + stage_sf_vectors = cute.zipped_divide(cute.slice_(stage_sf, (0, None)), (4,)) + for scatter_round in cutlass.range_constexpr(ceil_div(k_groups, 32)): + k_group_idx = Int32(scatter_round * 32) + lane_idx + if k_group_idx < Int32(k_groups): + cute.copy( + scatter_atom, + stage_sf_vectors[None, k_group_idx], + cute.zipped_divide( + sf_group[transfer_warp_idx, k_group_idx, None], + (4,), + )[None, Int32(token_in_group)], + ) + else: + pull_phase_mask = pull_phase_mask + cute.arch.sync_warp() + + # The activation rows have to be visible before the flags, and the flags carry release + # semantics that cover the scale stores, so one drain covers the whole group. + if lane_idx == Int32(0): + cute.arch.cp_async_bulk_wait_group(0) + cute.arch.sync_warp() + + sf_block_byte_base = Int64(expert_sf_begin + block_idx * Int32(128)) // Int64(128) * Int64( + sf_block_bytes + ) + Int64(column_in_block) * Int64(4 * group_size) + for store_round in cutlass.range_constexpr(ceil_div(k_groups, 32)): + k_group_idx = Int32(store_round * 32) + lane_idx + if k_group_idx < Int32(k_groups): + destination_sf = cute.make_tensor( + cute.make_ptr( + self.activation_sf_dtype, + fc1_activation_sf_pointer.toint() + + sf_block_byte_base + + Int64(k_group_idx) * Int64(512), + AddressSpace.gmem, + assumed_align=16, + ), + cute.make_layout(4 * group_size), + ) + cute.copy( + store_atom, + sf_group[transfer_warp_idx, k_group_idx, None], + destination_sf, + ) + cute.arch.sync_warp() + + for token_in_group in cutlass.range_constexpr(group_size): + token_in_expert = block_idx * Int32(128) + column_in_block + Int32(32 * token_in_group) + if token_in_expert < expert_token_count: + ready_slot_idx = expert_ready_slot_begin + token_in_expert // Int32( + self.tokens_per_fc1_ready_slot + ) + flag_tracker = flag_tracker.accumulate( + Int32(0), + self.token_in_flag_batch, + (fc1_ready_counter + ready_slot_idx).toint(), + ) + else: + flag_tracker = flag_tracker + cute.arch.sync_warp() + + if cutlass.const_expr(group_stage_sets == 1): + # The ring is exactly one group wide, so the stages are free only now. + if group_round + Int32(1) < group_count: + self._issue_token_group( + smem_workspace, + smem_base, + group_round + Int32(1), + first_group_in_this_expert, + expert_group_begin, + global_warp_count, + expert_token_count, + pool_expert_base, + next_stage_base, + transfer_warp_idx, + lane_idx, + ) + stage_base = next_stage_base + + first_group_in_this_expert = first_group_in_this_expert + group_count * global_warp_count + expert_group_begin = expert_group_end + expert_sf_begin = expert_sf_begin + ( + (expert_token_count + Int32(self.sf_padding_block - 1)) // Int32(self.sf_padding_block) + ) * Int32(self.sf_padding_block) + expert_ready_slot_begin = expert_ready_slot_begin + ( + (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) + // Int32(self.tokens_per_fc1_ready_slot) + ) + local_expert = local_expert + Int32(1) + + flag_tracker.fire() + cute.arch.sync_warp() + + @cute.jit + def _issue_token_group( + self, + smem_workspace, + smem_base, + group_round, + first_group_in_this_expert, + expert_group_begin, + worker_count, + expert_token_count, + pool_expert_base, + stage_base, + warp_slot, + lane_idx, + ) -> None: + """Start the remote reads for one token group into stages ``stage_base .. stage_base + G - 1``. + + A group is ``{j, j + 32, j + 64, j + 96}`` of one 128-token atom block, and the unit the worker space + strides over is the ``(block, j)`` pair rather than the token. Only the local staging address is + needed here, and that is pure arithmetic on the static schedule, so the lookahead costs no extra + dependent load; the outbound destination needs ``TokenSrcMetadata`` and stays on the drain side. + + A method rather than a closure because CuTe-DSL will not let a closure defined in one traced region + capture staged values, and that includes ``self``. + """ + group_size = self.dispatch_sf_group + token_in_activation = smem_workspace.tensor(self.token_in_activation_smem_region, smem_base) + token_in_sf = smem_workspace.tensor(self.token_in_sf_smem_region, smem_base) + warp_mbarrier = smem_workspace.ptr(self.token_in_mbarrier_region, smem_base) + warp_slot * Int32( + self.dispatch_stage_depth + ) + metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + activation_bytes = self.bytes_per_token + activation_sf_bytes = self.activation_sf_hidden_padded * int(self.activation_sf_dtype.width) // 8 + + group_in_this_expert = first_group_in_this_expert + group_round * worker_count - expert_group_begin + block_idx = group_in_this_expert // Int32(32) + column_in_block = group_in_this_expert - block_idx * Int32(32) + for token_in_group in cutlass.range_constexpr(group_size): + token_in_expert = block_idx * Int32(128) + column_in_block + Int32(32 * token_in_group) + if token_in_expert < expert_token_count: + pool_token_idx = pool_expert_base + token_in_expert + metadata = TokenSrcMetadata.load( + metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) + ) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), metadata.src_rank, Int64(0)) + remote_activation_address = ( + self._token_comm_args.activation.iterator.toint() + + peer_offset + + Int64(metadata.src_token) * Int64(self.bytes_per_token) + ) + remote_sf = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.activation_sf.dtype, + self._token_comm_args.activation_sf.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.activation_sf.iterator.max_alignment), + ), + self._token_comm_args.activation_sf.layout, + ) + remote_sf_row = remote_sf[Int64(metadata.src_token), None] + stage_idx = stage_base + Int32(token_in_group) + stage_mbarrier = warp_mbarrier + stage_idx + if lane_idx == Int32(0): + cute.arch.mbarrier_arrive_and_expect_tx( + stage_mbarrier, Int32(activation_bytes + activation_sf_bytes) + ) + tma_load_1d( + token_in_activation[warp_slot, stage_idx, None].iterator, + Int64(remote_activation_address), + stage_mbarrier, + Int32(activation_bytes), + ) + tma_load_1d( + token_in_sf[warp_slot, stage_idx, (None, None)].iterator, + remote_sf_row.iterator, + stage_mbarrier, + Int32(activation_sf_bytes), + ) + cute.arch.sync_warp() + @cute.jit def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: """Wait for pushed metadata, then pull activation payloads into local pools.""" - transfer_warp_idx = self._transfer_warp_idx + transfer_warp_idx = self._dispatch_warp_idx + dispatch_worker_idx = self._dispatch_worker_idx lane_idx = self._lane_idx - global_warp_idx = self._linear_cta_idx * Int32(self.transfer_warp_count) + transfer_warp_idx - global_warp_count = Int32(self.promised_launchable_sm_count * self.transfer_warp_count) + global_warp_idx = dispatch_worker_idx + global_warp_count = Int32(self.dispatch_cta_count * self.dispatch_warp_count) sizes = self._router.sizes_tensor(self._device_workspace) pool_expert_bases = self._router.pool_expert_base_tensor(self._device_workspace) @@ -1888,27 +2130,25 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) owner_expert_begin = self._local_rank * Int32(self.experts_per_rank) source_sizes = cute.make_tensor( - sizes.iterator + owner_expert_begin, cute.make_layout((self.experts_per_rank,)) + sizes.iterator + owner_expert_begin, + cute.make_layout((self.experts_per_rank,)), ) copy_elements = 4 if self.experts_per_rank % 4 == 0 else 1 source_size_vectors = cute.zipped_divide( - ( - mark_alignment(source_sizes, 16) - if cutlass.const_expr(copy_elements == 4) - else source_sizes - ), + (mark_alignment(source_sizes, 16) if cutlass.const_expr(copy_elements == 4) else source_sizes), (copy_elements,), ) destination_size_vectors = cute.zipped_divide(owned_sizes, (copy_elements,)) size_vector_count = cute.size(destination_size_vectors, mode=[1]) size_copy_atom = _copy_atom(cutlass.Int32, copy_elements * 32) - size_copy_rounds = ceil_div(size_vector_count, self.transfer_thread_count) - size_copy_registers = cute.make_rmem_tensor( - (copy_elements, size_copy_rounds), cutlass.Int32 - ) + # Stride by the role's own thread count, not the class constant: with a split dispatch owns + # twelve warps rather than the fused world's four, and the barrier below counts the same set. + dispatch_thread_count = self.dispatch_warp_count * 32 + size_copy_rounds = ceil_div(size_vector_count, dispatch_thread_count) + size_copy_registers = cute.make_rmem_tensor((copy_elements, size_copy_rounds), cutlass.Int32) transfer_thread_idx = transfer_warp_idx * Int32(32) + lane_idx for size_copy_round in cutlass.range_constexpr(size_copy_rounds): - vector_idx = Int32(size_copy_round * self.transfer_thread_count) + transfer_thread_idx + vector_idx = Int32(size_copy_round * dispatch_thread_count) + transfer_thread_idx if vector_idx < Int32(size_vector_count): cute.copy( size_copy_atom, @@ -1922,7 +2162,7 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No iket.range_pop() for size_copy_round in cutlass.range_constexpr(size_copy_rounds): - vector_idx = Int32(size_copy_round * self.transfer_thread_count) + transfer_thread_idx + vector_idx = Int32(size_copy_round * dispatch_thread_count) + transfer_thread_idx if vector_idx < Int32(size_vector_count): cute.copy( size_copy_atom, @@ -1931,18 +2171,31 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No ) iket.range_push("token_in.size_barrier") token_in_size_barrier = pipeline.NamedBarrier( - barrier_id=self.token_in_size_barrier_id, num_threads=self.transfer_thread_count + barrier_id=self.token_in_size_barrier_id, + num_threads=dispatch_thread_count, ) token_in_size_barrier.arrive_and_wait() iket.range_pop() - if cutlass.const_expr(self.token_back_mode == "standalone_warps"): # nosec B105 + if cutlass.const_expr(not self.cta_specialized and self.token_back_mode == "standalone_warps"): # nosec B105 sizes_ready_barrier = pipeline.NamedBarrier( barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count, ) sizes_ready_barrier.arrive() - iket.range_push("token_in.pull_payload") + if cutlass.const_expr(self.dispatch_stage_depth > 1): + self._token_in_multi_stage( + smem_workspace, + smem_base, + owned_sizes, + pool_expert_bases, + transfer_warp_idx, + lane_idx, + global_warp_idx, + global_warp_count, + ) + return + token_in_mbarriers = smem_workspace.ptr(self.token_in_mbarrier_region, smem_base) token_in_activation = smem_workspace.tensor(self.token_in_activation_smem_region, smem_base) token_in_sf = smem_workspace.tensor(self.token_in_sf_smem_region, smem_base) @@ -1956,15 +2209,14 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No fc1_activation_pointer = self._device_workspace.ptr(self.fc1_activation_region) fc1_activation_sf = self.fc1_activation_sf_tensor(self._device_workspace) fc1_ready_counter = self._device_workspace.ptr(self.fc1_ready_region) - activation_bytes = ( - cute.cosize(warp_activation_stage) * int(self.activation_dtype.width) // 8 - ) + activation_bytes = cute.cosize(warp_activation_stage) * int(self.activation_dtype.width) // 8 activation_sf_bytes = cute.cosize(warp_sf_stage) * int(self.activation_sf_dtype.width) // 8 sf_copy_elements = 4 source_sf_values = cute.slice_(warp_sf_stage, (0, None)) source_sf_vectors = cute.zipped_divide(source_sf_values, (sf_copy_elements,)) sf_copy_atom = _copy_atom( - self.activation_sf_dtype, sf_copy_elements * int(self.activation_sf_dtype.width) + self.activation_sf_dtype, + sf_copy_elements * int(self.activation_sf_dtype.width), ) next_dense_token = global_warp_idx @@ -1986,26 +2238,19 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No expert_valid_end = expert_valid_begin + expert_token_count pull_count = Int32(0) if next_dense_token < expert_valid_end: - pull_count = ( - expert_valid_end - next_dense_token + global_warp_count - Int32(1) - ) // global_warp_count + pull_count = (expert_valid_end - next_dense_token + global_warp_count - Int32(1)) // global_warp_count for pull_round in cutlass.range(pull_count, unroll=1): dense_token_idx = next_dense_token + Int32(pull_round) * global_warp_count token_in_expert = dense_token_idx - expert_valid_begin pool_token_idx = pool_expert_bases[local_expert] + token_in_expert sf_token_idx = expert_sf_begin + token_in_expert - ready_slot_idx = expert_ready_slot_begin + token_in_expert // Int32( - self.tokens_per_fc1_ready_slot - ) + ready_slot_idx = expert_ready_slot_begin + token_in_expert // Int32(self.tokens_per_fc1_ready_slot) metadata = TokenSrcMetadata.load( - token_metadata_pointer.toint() - + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) - ) - peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map( - Int64(0), metadata.src_rank, Int64(0) + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) ) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), metadata.src_rank, Int64(0)) remote_activation_address = ( self._token_comm_args.activation.iterator.toint() + peer_offset @@ -2043,8 +2288,7 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No destination_activation = cute.make_ptr( self.activation_dtype, - fc1_activation_pointer.toint() - + Int64(pool_token_idx) * Int64(self.bytes_per_token), + fc1_activation_pointer.toint() + Int64(pool_token_idx) * Int64(self.bytes_per_token), AddressSpace.gmem, assumed_align=16, ) @@ -2060,9 +2304,7 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No destination_sf_row = fc1_activation_sf[Int64(sf_token_idx), ((None, None), None)] destination_sf_values = cute.slice_(destination_sf_row, (0, None, None)) destination_sf_values = cute.group_modes(destination_sf_values, 0, 2) - destination_sf_vectors = cute.zipped_divide( - destination_sf_values, (sf_copy_elements,) - ) + destination_sf_vectors = cute.zipped_divide(destination_sf_values, (sf_copy_elements,)) sf_vector_count = cute.size(destination_sf_vectors, mode=[1]) for sf_round in cutlass.range_constexpr(ceil_div(sf_vector_count, 32)): sf_vector_idx = Int32(sf_round * 32) + lane_idx @@ -2076,17 +2318,14 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No cute.arch.cp_async_bulk_wait_group(0, read=True) cute.arch.sync_warp() ready_address = (fc1_ready_counter + ready_slot_idx).toint() - flag_tracker = flag_tracker.accumulate( - Int32(0), self.token_in_flag_batch, ready_address - ) + flag_tracker = flag_tracker.accumulate(Int32(0), self.token_in_flag_batch, ready_address) cute.arch.sync_warp() pull_phase = pull_phase ^ Int32(1) next_dense_token = next_dense_token + pull_count * global_warp_count expert_valid_begin = expert_valid_end expert_sf_begin = expert_sf_begin + ( - (expert_token_count + Int32(self.sf_padding_block - 1)) - // Int32(self.sf_padding_block) + (expert_token_count + Int32(self.sf_padding_block - 1)) // Int32(self.sf_padding_block) ) * Int32(self.sf_padding_block) expert_ready_slot_begin = expert_ready_slot_begin + ( (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) @@ -2096,13 +2335,11 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No flag_tracker.fire() cute.arch.sync_warp() - iket.range_pop() - if cutlass.const_expr( - self.token_back_enabled and self.token_back_mode != "standalone_warps" # nosec B105 - ): + if cutlass.const_expr(self.token_back_enabled and self.token_back_mode != "standalone_warps"): # nosec B105 iket.range_push("token_in.transfer_barrier") transfer_lifetime_barrier = pipeline.NamedBarrier( - barrier_id=self.transfer_lifetime_barrier_id, num_threads=self.transfer_thread_count + barrier_id=self.transfer_lifetime_barrier_id, + num_threads=dispatch_thread_count, ) transfer_lifetime_barrier.arrive_and_wait() iket.range_pop() @@ -2114,31 +2351,31 @@ def _stateless_pace(self, reference_window: Int32, current_window: Int32) -> Non sleep_cycles = reference_window - current_window elif current_window > reference_window: sleep_cycles = cutlass.min( - current_window - reference_window, Int32(self.standalone_max_backoff_cycles) + current_window - reference_window, + Int32(self.standalone_max_backoff_cycles), ) if sleep_cycles > Int32(0): nanosleep(sleep_cycles) @cute.jit def _adaptive_pace( - self, average_window: Int32, current_window: Int32, low_window: int, high_window: int + self, + average_window: Int32, + current_window: Int32, + low_window: int, + high_window: int, ) -> Int32: sleep_cycles = Int32(0) if current_window > average_window: sleep_cycles = current_window - average_window - average_window = average_window + ( - (current_window - average_window + Int32(3)) // Int32(4) - ) + average_window = average_window + ((current_window - average_window + Int32(3)) // Int32(4)) if sleep_cycles > Int32(high_window): sleep_cycles = Int32(high_window) else: sleep_cycles = average_window - current_window - average_window = average_window - ( - (average_window - current_window + Int32(3)) // Int32(4) - ) + average_window = average_window - ((average_window - current_window + Int32(3)) // Int32(4)) # with cute.arch.elect_one(): - # cute.printf("avg_window: {}, current_window: {}, sleep_cycles: {}", average_window, current_window, - # sleep_cycles) + # cute.printf("avg_window: {}, current_window: {}, sleep_cycles: {}", average_window, current_window, sleep_cycles) if sleep_cycles > Int32(self.adaptive_minimum_sleep_cycles): nanosleep(sleep_cycles) if average_window > Int32(high_window): @@ -2147,29 +2384,269 @@ def _adaptive_pace( average_window = Int32(low_window) return average_window + @cute.jit + def _issue_combine_read( + self, + smem_workspace, + smem_base, + chunk_round, + stage_index, + first_dense_token, + expert_valid_begin, + pool_expert_base, + worker_count, + warp_slot, + lane_idx, + ) -> None: + """Read chunk ``chunk_round`` of this worker's stream from the local FC2 staging into ``stage_index``. + + Needs no metadata: the local address is arithmetic on the static schedule's ``first + (r // C) * W``, + so running the read ahead costs one integer divide and no extra dependent load. It is the OUTBOUND + destination that needs ``TokenSrcMetadata``, and that stays on the drain side where the wait for this + read has already been paid. + """ + chunk_bytes = self.token_back_chunk_bytes + chunk_count = ceil_div(self.combine_output_token_bytes, chunk_bytes) + stage = smem_workspace.tensor(self.token_back_activation_smem_region, smem_base) + stage_mbarrier = ( + smem_workspace.ptr(self.token_back_mbarrier_region, smem_base) + + warp_slot * Int32(self.combine_stage_depth) + + stage_index + ) + fc2_activation_pointer = self._device_workspace.ptr(self.fc2_activation_region) + + token_round = chunk_round // Int32(chunk_count) + chunk_idx = chunk_round - token_round * Int32(chunk_count) + dense_token_idx = first_dense_token + token_round * worker_count + pool_token_idx = pool_expert_base + (dense_token_idx - expert_valid_begin) + current_chunk_bytes = Int32(chunk_bytes) + remaining_bytes = Int32(self.combine_output_token_bytes) - chunk_idx * Int32(chunk_bytes) + if remaining_bytes < current_chunk_bytes: + current_chunk_bytes = remaining_bytes + else: + current_chunk_bytes = current_chunk_bytes + if lane_idx == Int32(0): + cute.arch.mbarrier_arrive_and_expect_tx(stage_mbarrier, current_chunk_bytes) + tma_load_1d( + stage[warp_slot, stage_index, None].iterator, + fc2_activation_pointer.toint() + + Int64(pool_token_idx) * Int64(self.combine_output_token_bytes) + + Int64(chunk_idx) * Int64(chunk_bytes), + stage_mbarrier, + current_chunk_bytes, + ) + + @cute.jit + def _token_back_multi_stage( + self, + smem_workspace, + smem_base, + owned_sizes, + transfer_warp_idx, + lane_idx, + global_worker_idx, + global_worker_count, + ) -> None: + """Push output chunks with the local read running the whole ring ahead of the outbound store.""" + depth = self.combine_stage_depth + chunk_bytes = self.token_back_chunk_bytes + chunk_count = ceil_div(self.combine_output_token_bytes, chunk_bytes) + + stage = smem_workspace.tensor(self.token_back_activation_smem_region, smem_base) + worker_mbarrier = smem_workspace.ptr(self.token_back_mbarrier_region, smem_base) + transfer_warp_idx * Int32( + depth + ) + if lane_idx == Int32(0): + for init_stage in cutlass.range_constexpr(depth): + cute.arch.mbarrier_init(worker_mbarrier + Int32(init_stage), 1) + cute.arch.sync_warp() + + if cutlass.const_expr(not self.cta_specialized and self.token_back_mode == "standalone_warps"): # nosec B105 + sizes_ready_barrier = pipeline.NamedBarrier( + barrier_id=self.standalone_size_barrier_id, + num_threads=2 * self.transfer_thread_count, + ) + sizes_ready_barrier.arrive_and_wait() + + pool_expert_bases = self._router.pool_expert_base_tensor(self._device_workspace) + token_metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + fc2_done = self._device_workspace.ptr(self.fc2_done_region) + + next_dense_token = global_worker_idx + push_phase_mask = Int32(0) + expert_valid_begin = Int32(0) + local_expert = Int32(0) + while local_expert < Int32(self.experts_per_rank): + expert_token_count = owned_sizes[local_expert] + expert_valid_end = expert_valid_begin + expert_token_count + token_count = Int32(0) + if next_dense_token < expert_valid_end: + token_count = ( + expert_valid_end - next_dense_token + global_worker_count - Int32(1) + ) // global_worker_count + token_tile_count = (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) // Int32( + self.tokens_per_fc1_ready_slot + ) + completion_target = token_tile_count * Int32(self.fc2_done_signals_per_token_tile) + iket.range_push("token_back.wait_fc2") + while cute.arch.load(fc2_done + local_expert, Int32, sem="acquire", scope="gpu") < completion_target: + nanosleep(500) + iket.range_pop() + chunk_total = token_count * Int32(chunk_count) + pool_expert_base = pool_expert_bases[local_expert] + + # The prologue issues with no recycle in front of it, so its stages must be free of the previous + # expert's stores as well as its own. One full read drain buys that, once per expert against + # thousands of chunks. + if chunk_total > Int32(0): + if lane_idx == Int32(0): + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.sync_warp() + for prologue_chunk in cutlass.range_constexpr(depth): + if Int32(prologue_chunk) < chunk_total: + self._issue_combine_read( + smem_workspace, + smem_base, + Int32(prologue_chunk), + Int32(prologue_chunk), + next_dense_token, + expert_valid_begin, + pool_expert_base, + global_worker_count, + transfer_warp_idx, + lane_idx, + ) + cute.arch.sync_warp() + + stage_idx = Int32(0) + for chunk_round in cutlass.range(chunk_total, unroll=1): + token_round = chunk_round // Int32(chunk_count) + chunk_idx = chunk_round - token_round * Int32(chunk_count) + dense_token_idx = next_dense_token + token_round * global_worker_count + pool_token_idx = pool_expert_base + (dense_token_idx - expert_valid_begin) + metadata = TokenSrcMetadata.load( + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) + ) + destination_topk = metadata.src_topk + if cutlass.const_expr(self.reduce_topk_in_kernel): + destination_topk = Int32(0) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), metadata.src_rank, Int64(0)) + remote_activation = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.pre_reduced_activation.dtype, + self._token_comm_args.pre_reduced_activation.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.pre_reduced_activation.iterator.max_alignment), + ), + self._token_comm_args.pre_reduced_activation.layout, + ) + destination_row = remote_activation[Int64(metadata.src_token), destination_topk, None] + current_chunk_bytes = Int32(chunk_bytes) + remaining_bytes = Int32(self.combine_output_token_bytes) - chunk_idx * Int32(chunk_bytes) + if remaining_bytes < current_chunk_bytes: + current_chunk_bytes = remaining_bytes + else: + current_chunk_bytes = current_chunk_bytes + destination_chunk = cute.make_ptr( + cutlass.Uint8, + destination_row.iterator.toint() + Int64(chunk_idx) * Int64(chunk_bytes), + AddressSpace.gmem, + assumed_align=16, + ) + cute.arch.mbarrier_wait( + worker_mbarrier + stage_idx, + (push_phase_mask >> stage_idx) & Int32(1), + ) + push_phase_mask = push_phase_mask ^ (Int32(1) << stage_idx) + if lane_idx == Int32(0): + if cutlass.const_expr(self.reduce_topk_in_kernel): + cp_reduce_async_bulk_add_bf16_s2g( + destination_chunk, + stage[transfer_warp_idx, stage_idx, None].iterator, + current_chunk_bytes, + ) + else: + cp_async_bulk_s2g( + destination_chunk, + stage[transfer_warp_idx, stage_idx, None].iterator, + current_chunk_bytes, + ) + cute.arch.cp_async_bulk_commit_group() + cute.arch.sync_warp() + + refill_round = chunk_round + Int32(depth) + if refill_round < chunk_total: + # Every stage is given to the reads, so the stage being refilled is the one whose store + # was issued `depth` chunks ago and nothing newer has to be complete. + if lane_idx == Int32(0): + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.sync_warp() + self._issue_combine_read( + smem_workspace, + smem_base, + refill_round, + stage_idx, + next_dense_token, + expert_valid_begin, + pool_expert_base, + global_worker_count, + transfer_warp_idx, + lane_idx, + ) + cute.arch.sync_warp() + stage_idx = stage_idx + Int32(1) + if stage_idx == Int32(depth): + stage_idx = Int32(0) + + next_dense_token = next_dense_token + token_count * global_worker_count + expert_valid_begin = expert_valid_end + local_expert = local_expert + Int32(1) + + # token_back publishes no flag, but the tail NvlinkBarrier's fence does not drain async-proxy bulk + # copies, so without this a peer could read a half-written row after the rendezvous. + if lane_idx == Int32(0): + cute.arch.cp_async_bulk_wait_group(0) + cute.arch.sync_warp() + @cute.jit def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: """Push completed FC2 data and scale rows to source ranks.""" - transfer_warp_idx = self._transfer_warp_idx + transfer_warp_idx = self._combine_warp_idx + combine_worker_idx = self._combine_worker_idx lane_idx = self._lane_idx if cutlass.const_expr(not self.token_back_enabled): return if cutlass.const_expr( - self.combine_format.is_quantized - and self._token_comm_args.pre_reduced_activation_sf is None + self.combine_format.is_quantized and self._token_comm_args.pre_reduced_activation_sf is None ): raise ValueError("Quantized token-back requires a scale destination.") - global_worker_idx = ( - self._linear_cta_idx * Int32(self.transfer_warp_count) + transfer_warp_idx - ) - global_worker_count = Int32(self.promised_launchable_sm_count * self.transfer_warp_count) + if cutlass.const_expr(self.cta_specialized): + # A dedicated token-back CTA need not have token-in's SMEM scratch. + # pool_expert_base is built after sizes_ready; acquire metadata_ready too. + self._router.wait_for_sizes_ready(self._device_workspace) + owned_sizes = self._router.local_expert_sizes(self._device_workspace, self._local_rank) + self._router.wait_for_metadata_ready(self._device_workspace) + else: + owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) + + global_worker_idx = combine_worker_idx + global_worker_count = Int32(self.combine_cta_count * self.combine_warp_count) + if cutlass.const_expr(self.combine_stage_depth > 1): + self._token_back_multi_stage( + smem_workspace, + smem_base, + owned_sizes, + transfer_warp_idx, + lane_idx, + global_worker_idx, + global_worker_count, + ) + return token_back_mbarriers = smem_workspace.ptr(self.token_back_mbarrier_region, smem_base) worker_mbarrier = token_back_mbarriers + transfer_warp_idx if cutlass.const_expr(self.token_back_push_data): - token_back_activation = smem_workspace.tensor( - self.token_back_activation_smem_region, smem_base - ) + token_back_activation = smem_workspace.tensor(self.token_back_activation_smem_region, smem_base) worker_activation_stage = token_back_activation[transfer_warp_idx, None] activation_chunk_bytes = ( cute.cosize(worker_activation_stage) * int(self.combine_format.act_dtype.width) // 8 @@ -2177,15 +2654,12 @@ def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> if cutlass.const_expr(self.token_back_push_sf): token_back_sf = smem_workspace.tensor(self.token_back_sf_smem_region, smem_base) worker_sf_stage = token_back_sf[transfer_warp_idx, (None, None)] - sf_chunk_bytes = ( - cute.cosize(worker_sf_stage) * int(self.combine_format.scale_dtype.width) // 8 - ) + sf_chunk_bytes = cute.cosize(worker_sf_stage) * int(self.combine_format.scale_dtype.width) // 8 if lane_idx == Int32(0): cute.arch.mbarrier_init(worker_mbarrier, 1) cute.arch.sync_warp() - owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) - if cutlass.const_expr(self.token_back_mode == "standalone_warps"): # nosec B105 + if cutlass.const_expr(not self.cta_specialized and self.token_back_mode == "standalone_warps"): # nosec B105 sizes_ready_barrier = pipeline.NamedBarrier( barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count, @@ -2201,8 +2675,7 @@ def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> activation_chunk_count = ceil_div(output_token_bytes, activation_chunk_bytes) data_window_unit = activation_chunk_bytes * 2 reuse_data_pacing_enabled = ( - self.token_back_mode == "reuse_dispatch_warps" # nosec B105 - and data_window_unit > self.minimum_pacing_window_cycles + self.token_back_mode == "reuse_dispatch_warps" and data_window_unit > self.minimum_pacing_window_cycles # nosec B105 ) # Preserve the empirical low:initial:high ratio of 1:2.5:5. data_average_window = Int32(data_window_unit) @@ -2210,14 +2683,11 @@ def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> data_high_window = data_window_unit * 2 if cutlass.const_expr(self.token_back_push_sf): fc2_sf_pointer = self._device_workspace.ptr(self.fc2_activation_sf_region) - output_sf_bytes = ( - self.combine_sf_hidden_padded * int(self.combine_format.scale_dtype.width) // 8 - ) + output_sf_bytes = self.combine_sf_hidden_padded * int(self.combine_format.scale_dtype.width) // 8 sf_chunk_count = ceil_div(output_sf_bytes, sf_chunk_bytes) sf_window_unit = ceil_div(sf_chunk_bytes * 2, 3) reuse_sf_pacing_enabled = ( - self.token_back_mode == "reuse_dispatch_warps" # nosec B105 - and sf_window_unit > self.minimum_pacing_window_cycles + self.token_back_mode == "reuse_dispatch_warps" and sf_window_unit > self.minimum_pacing_window_cycles # nosec B105 ) sf_average_window = Int32(sf_window_unit) sf_low_window = sf_window_unit * 2 // 5 @@ -2229,60 +2699,63 @@ def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> next_dense_token = self.next_token(next_dense_token) expert_valid_begin = Int32(0) transfer_phase = Int32(0) + if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): # nosec B105 + expert_ready_slot_begin = Int32(0) + last_ready_slot = Int32(-1) - iket.range_push("token_back.work") local_expert = Int32(0) while local_expert < Int32(self.experts_per_rank): expert_token_count = owned_sizes[local_expert] expert_valid_end = expert_valid_begin + expert_token_count - if next_dense_token < expert_valid_end: - token_tile_count = ( - expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1) - ) // Int32(self.tokens_per_fc1_ready_slot) - completion_target = token_tile_count * Int32(self.fc2_done_signals_per_token_tile) - iket.range_push("token_back.wait_fc2") - while ( - cute.arch.load(fc2_done + local_expert, Int32, sem="acquire", scope="gpu") - < completion_target - ): - nanosleep(500) - iket.range_pop() + if cutlass.const_expr(self.token_back_ready_granularity == "expert"): # nosec B105 + if next_dense_token < expert_valid_end: + token_tile_count = (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) // Int32( + self.tokens_per_fc1_ready_slot + ) + completion_target = token_tile_count * Int32(self.fc2_done_signals_per_token_tile) + iket.range_push("token_back.wait_fc2") + while cute.arch.load(fc2_done + local_expert, Int32, sem="acquire", scope="gpu") < completion_target: + nanosleep(500) + iket.range_pop() while next_dense_token < expert_valid_end: token_in_expert = next_dense_token - expert_valid_begin + if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): # nosec B105 + # Same rank-local expert/tile prefix as FC1 ready; separate counter storage. + ready_slot = expert_ready_slot_begin + token_in_expert // Int32(self.tokens_per_fc1_ready_slot) + if ready_slot != last_ready_slot: + iket.range_push("token_back.wait_fc2") + while cute.arch.load(fc2_done + ready_slot, Int32, sem="acquire", scope="gpu") < Int32( + self.fc2_done_signals_per_token_tile + ): + nanosleep(500) + iket.range_pop() + last_ready_slot = ready_slot pool_token_idx = pool_expert_bases[local_expert] + token_in_expert metadata = TokenSrcMetadata.load( - token_metadata_pointer.toint() - + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) ) destination_topk = metadata.src_topk if cutlass.const_expr(self.reduce_topk_in_kernel): destination_topk = Int32(0) - peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map( - Int64(0), metadata.src_rank, Int64(0) - ) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), metadata.src_rank, Int64(0)) is_remote_token = metadata.src_rank != self._local_rank if cutlass.const_expr(self.token_back_push_data): iket.range_push("token_back.push_data") - local_activation_address = fc2_activation_pointer.toint() + Int64( - pool_token_idx - ) * Int64(output_token_bytes) + local_activation_address = fc2_activation_pointer.toint() + Int64(pool_token_idx) * Int64( + output_token_bytes + ) remote_activation = cute.make_tensor( cute.make_ptr( self._token_comm_args.pre_reduced_activation.dtype, - self._token_comm_args.pre_reduced_activation.iterator.toint() - + peer_offset, + self._token_comm_args.pre_reduced_activation.iterator.toint() + peer_offset, AddressSpace.gmem, - assumed_align=( - self._token_comm_args.pre_reduced_activation.iterator.max_alignment - ), + assumed_align=(self._token_comm_args.pre_reduced_activation.iterator.max_alignment), ), self._token_comm_args.pre_reduced_activation.layout, ) - destination_row = remote_activation[ - Int64(metadata.src_token), destination_topk, None - ] + destination_row = remote_activation[Int64(metadata.src_token), destination_topk, None] for chunk_idx in cutlass.range_constexpr(activation_chunk_count): chunk_byte_offset = Int64(chunk_idx * activation_chunk_bytes) chunk_bytes_this_round = min( @@ -2297,17 +2770,13 @@ def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> ) round_start_clock = Int64(0) cute.arch.sync_warp() - if cutlass.const_expr( - reuse_data_pacing_enabled or stateless_data_pacing_enabled - ): + if cutlass.const_expr(reuse_data_pacing_enabled or stateless_data_pacing_enabled): if is_remote_token: round_start_clock = read_clock64() else: round_start_clock = round_start_clock with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - worker_mbarrier, current_chunk_bytes - ) + cute.arch.mbarrier_arrive_and_expect_tx(worker_mbarrier, current_chunk_bytes) tma_load_1d( worker_activation_stage.iterator, local_activation_address + chunk_byte_offset, @@ -2355,29 +2824,20 @@ def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> if cutlass.const_expr(self.token_back_push_sf): iket.range_push("token_back.push_sf") - local_sf_address = fc2_sf_pointer.toint() + Int64(pool_token_idx) * Int64( - output_sf_bytes - ) + local_sf_address = fc2_sf_pointer.toint() + Int64(pool_token_idx) * Int64(output_sf_bytes) remote_sf = cute.make_tensor( cute.make_ptr( self._token_comm_args.pre_reduced_activation_sf.dtype, - self._token_comm_args.pre_reduced_activation_sf.iterator.toint() - + peer_offset, + self._token_comm_args.pre_reduced_activation_sf.iterator.toint() + peer_offset, AddressSpace.gmem, - assumed_align=( - self._token_comm_args.pre_reduced_activation_sf.iterator.max_alignment - ), + assumed_align=(self._token_comm_args.pre_reduced_activation_sf.iterator.max_alignment), ), self._token_comm_args.pre_reduced_activation_sf.layout, ) - destination_sf_row = remote_sf[ - Int64(metadata.src_token), destination_topk, None - ] + destination_sf_row = remote_sf[Int64(metadata.src_token), destination_topk, None] for chunk_idx in cutlass.range_constexpr(sf_chunk_count): chunk_byte_offset = Int64(chunk_idx * sf_chunk_bytes) - chunk_bytes_this_round = min( - sf_chunk_bytes, output_sf_bytes - chunk_idx * sf_chunk_bytes - ) + chunk_bytes_this_round = min(sf_chunk_bytes, output_sf_bytes - chunk_idx * sf_chunk_bytes) current_chunk_bytes = Int32(chunk_bytes_this_round) current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) stateless_sf_pacing_enabled = ( @@ -2386,17 +2846,13 @@ def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> ) round_start_clock = Int64(0) cute.arch.sync_warp() - if cutlass.const_expr( - reuse_sf_pacing_enabled or stateless_sf_pacing_enabled - ): + if cutlass.const_expr(reuse_sf_pacing_enabled or stateless_sf_pacing_enabled): if is_remote_token: round_start_clock = read_clock64() else: round_start_clock = round_start_clock with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - worker_mbarrier, current_chunk_bytes - ) + cute.arch.mbarrier_arrive_and_expect_tx(worker_mbarrier, current_chunk_bytes) tma_load_1d( worker_sf_stage.iterator, local_sf_address + chunk_byte_offset, @@ -2412,7 +2868,9 @@ def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> ) with cute.arch.elect_one(): cp_async_bulk_s2g( - destination_chunk, worker_sf_stage.iterator, current_chunk_bytes + destination_chunk, + worker_sf_stage.iterator, + current_chunk_bytes, ) cute.arch.cp_async_bulk_commit_group() cute.arch.cp_async_bulk_wait_group(0) @@ -2422,7 +2880,10 @@ def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> if is_remote_token: current_window = Int32(read_clock64() - round_start_clock) sf_average_window = self._adaptive_pace( - sf_average_window, current_window, sf_low_window, sf_high_window + sf_average_window, + current_window, + sf_low_window, + sf_high_window, ) elif cutlass.const_expr(stateless_sf_pacing_enabled): if is_remote_token: @@ -2433,24 +2894,28 @@ def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> cute.arch.sync_warp() next_dense_token = self.next_token(next_dense_token) + if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): # nosec B105 + # Advance even when this worker skipped the expert (empty experts contribute zero). + expert_ready_slot_begin = expert_ready_slot_begin + ( + expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1) + ) // Int32(self.tokens_per_fc1_ready_slot) expert_valid_begin = expert_valid_end local_expert = local_expert + Int32(1) - iket.range_pop() - # with cute.arch.elect_one(): - # cute.printf(" final data_average_window: {} ", data_average_window) @cute.jit def next_token(self, current_token: Int32) -> Int32: - global_worker_count = self.promised_launchable_sm_count * self.transfer_warp_count + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + # The static stride is the number of warps that actually walk the pool. With a split that is the + # communication CTAs' combine warps, not every persistent CTA's transfer warps. + global_worker_count = self.combine_cta_count * self.combine_warp_count schedule_counter = None if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 schedule_counter = self._device_workspace.ptr(self.token_back_schedule_region) if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 claimed_token = Int32(0) - if self._lane_idx == Int32(0): - claimed_token = cute.arch.atomic_add( - schedule_counter, Int32(1), sem="relaxed", scope="gpu" - ) + if lane_idx == Int32(0): + claimed_token = cute.arch.atomic_add(schedule_counter, Int32(1), sem="relaxed", scope="gpu") return Int32(cute.arch.shuffle_sync(claimed_token, Int32(0))) return current_token + global_worker_count @@ -2471,9 +2936,7 @@ def reset_tail(self) -> None: ) iket.range_pop() total_reset_threads = self.promised_launchable_sm_count * self.transfer_thread_count - global_reset_thread = ( - self._linear_cta_idx * Int32(self.transfer_thread_count) + transfer_thread_idx - ) + global_reset_thread = self._linear_cta_idx * Int32(self.transfer_thread_count) + transfer_thread_idx iket.range_push("tail.reset_workspace") self._device_workspace.reset_tail_space("shared", global_reset_thread, total_reset_threads) self._device_workspace.reset_tail_space("local", global_reset_thread, total_reset_threads) @@ -2500,4 +2963,4 @@ def reset_tail(self) -> None: iket.range_pop() -__all__ = ["TokenBackMode", "TokenBackScheduleMode", "TokenCommArgs", "TokenCommNonDeterministic"] +__all__ = ["TokenCommNonDeterministic", "TokenCommArgs"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_cta_mainloop.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_cta_mainloop.py new file mode 100644 index 000000000000..ecc713fcbf9d --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_cta_mainloop.py @@ -0,0 +1,93 @@ +"""Warp execution and scratch storage for dedicated token communication CTAs.""" + +import cutlass +import cutlass.cute as cute +from cutlass.cutlass_dsl import Int32 + +from ...helpers.smem_workspace import SmemWorkspace +from .token_comm import TokenCommNonDeterministic + + +class TokenCommCtaMainloop: + """Run token-in and optional concurrent token-back on a dedicated CTA. + + Device buffers and counters belong to the shared TokenComm protocol. This + mainloop owns the communication CTA's warp branches and SMEM placement; + each transfer operation initializes and drains its own asynchronous copies. + The enclosing kernel performs the all-CTA reset after both sides finish. + """ + + def __init__( + self, + token_comm: TokenCommNonDeterministic, + *, + token_back_on_compute_ctas: bool, + ) -> None: + self.token_comm = token_comm + self.token_back_on_compute_ctas = token_back_on_compute_ctas + self.dispatch_warp_start = token_comm.dispatch_warp_start + self.dispatch_warp_count = token_comm.dispatch_warp_count + self.combine_warp_start = token_comm.combine_warp_start + self.combine_warp_count = token_comm.combine_warp_count + if token_comm.token_back_mode != "standalone_warps": # nosec B105 + raise ValueError("Dedicated communication CTAs require standalone token-back warps.") + if not token_back_on_compute_ctas: + dispatch_end = self.dispatch_warp_start + self.dispatch_warp_count + combine_end = self.combine_warp_start + self.combine_warp_count + if self.dispatch_warp_start < combine_end and self.combine_warp_start < dispatch_end: + raise ValueError("Concurrent token-in and token-back warp intervals must not overlap.") + + def register_smem_regions(self, workspace: SmemWorkspace) -> None: + """Declare communication CTA scratch independently of fused placement.""" + self.register_token_in_smem(workspace) + if not self.token_back_on_compute_ctas: + self.register_token_back_smem(workspace) + + def register_token_in_smem(self, workspace: SmemWorkspace) -> None: + """Declare input scratch for the dedicated dispatch warp group.""" + comm = self.token_comm + workspace.register_mbarrier( + comm.token_in_mbarrier_region, + self.dispatch_warp_count * comm.dispatch_stage_depth, + ) + workspace.register_tensor( + comm.expert_sizes_smem_region, + cutlass.Int32, + (comm.experts_per_rank,), + byte_alignment=16, + ) + staging = workspace.create_overlay("nvlink.token_comm.main_smem.transfer").add_lifetime("token_in") + comm._register_token_in_staging(staging) + + def register_token_back_smem(self, workspace: SmemWorkspace) -> None: + """Declare output scratch on whichever CTA executes standalone token-back.""" + comm = self.token_comm + if not comm.token_back_enabled: + return + workspace.register_mbarrier( + comm.token_back_mbarrier_region, + self.combine_warp_count * comm.combine_stage_depth, + ) + staging = workspace.create_overlay("nvlink.token_comm.main_smem.standalone_token_back").add_lifetime( + "token_back" + ) + comm._register_token_back_staging(staging) + + @cute.jit + def run( + self, + smem_workspace: SmemWorkspace, + smem_base: cute.Pointer, + ) -> None: + """Run independent warp groups without joining them before their drains.""" + thread_idx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.make_warp_uniform(thread_idx // Int32(32)) + if (warp_idx >= Int32(self.dispatch_warp_start)) & ( + warp_idx < Int32(self.dispatch_warp_start + self.dispatch_warp_count) + ): + TokenCommNonDeterministic.token_in(self.token_comm, smem_workspace, smem_base) + elif cutlass.const_expr(not self.token_back_on_compute_ctas): + if (warp_idx >= Int32(self.combine_warp_start)) & ( + warp_idx < Int32(self.combine_warp_start + self.combine_warp_count) + ): + TokenCommNonDeterministic.token_back(self.token_comm, smem_workspace, smem_base) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_deterministic.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_deterministic.py new file mode 100644 index 000000000000..af3e8b6545b2 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_deterministic.py @@ -0,0 +1,2576 @@ +"""Metadata-push routing with a fixed token sequence and fused communication.""" + +import dataclasses +import os +from typing import Callable, ClassVar, Optional, Tuple, Union + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils +from cutlass._mlir.dialects import llvm +from cutlass.cute.typing import AddressSpace +from cutlass.cutlass_dsl import Int32, Int64 +from cutlass.utils.blockscaled_layout import tile_atom_to_shape_SF + +from ...api import ImplDesc, KernelComponent, OptionalRequirement, ProblemDesc +from ...helpers.device_workspace import DeviceWorkspace +from ...helpers.dsl_helpers import smem_exclusive_prefix +from ...helpers.flag_batch import make_flag_batch_tracker +from ...helpers.iket_compat import iket +from ...helpers.software_sync import NvlinkBarrier +from ...helpers.ptx_helpers import ( + cp_async_bulk_s2g, + cp_reduce_async_bulk_add_bf16_s2g, + cp_reduce_async_bulk_add_u32_s2g, + nanosleep, + read_clock64, + red_add_relaxed_sys_s32, + stg_b64, + stg_f32, + tma_load_1d, +) +from ...helpers.smem_workspace import SmemWorkspace +from ...helpers.utils import ceil_div, round_up +from ...quant_def import CombineFormat, QuantKind +from ..token_protocol import TokenSrcMetadata +from .symmetric_buffer import SymmetricBufferDevice +from .token_comm import TokenBackMode, TokenBackScheduleMode, TokenCommArgs + + +_quant_spec = { + "nvfp4": (cutlass.Float4E2M1FN, cutlass.Float8E4M3FN, 16), + "mxfp4": (cutlass.Float4E2M1FN, cutlass.Float8E8M0FNU, 32), + "mxfp8_e4m3": (cutlass.Float8E4M3FN, cutlass.Float8E8M0FNU, 32), + "mxfp8_e5m2": (cutlass.Float8E5M2, cutlass.Float8E8M0FNU, 32), +} + + +@dataclasses.dataclass(frozen=True) +class _ReceiveCapacity: + raw_route_count: int + logical_route_count: int + padded_route_count: int + + +def _compute_receive_capacity( + *, + world_size: int, + max_tokens_per_rank: int, + topk: int, + experts_per_rank: int, + max_recv_size_per_rank: int, + padding_block: int, +) -> _ReceiveCapacity: + raw_route_count = world_size * max_tokens_per_rank * topk + logical_route_count = min(max_recv_size_per_rank, raw_route_count) + active_expert_count = min(experts_per_rank, logical_route_count) + padded_block_count = active_expert_count + (logical_route_count - active_expert_count) // padding_block + return _ReceiveCapacity( + raw_route_count=raw_route_count, + logical_route_count=logical_route_count, + padded_route_count=padded_block_count * padding_block, + ) + + +@dataclasses.dataclass(frozen=True) +class _SortedElement: + flat_topk_index: Int32 + topk_score: Optional[cutlass.Float32] + + def pack(self) -> Union[Int64, Int32]: + if cutlass.const_expr(self.topk_score is None): + return self.flat_topk_index + scratch = cute.make_rmem_tensor((2,), cutlass.Int32) + scratch[0] = self.flat_topk_index + cute.recast_tensor(scratch, cutlass.Float32)[1] = self.topk_score + return cute.recast_tensor(scratch, cutlass.Int64)[0] + + @classmethod + def from_packed(cls, packed: Union[Int64, Int32]) -> "_SortedElement": + if cutlass.const_expr(type(packed).width == 32): + return cls(flat_topk_index=Int32(packed), topk_score=None) + scratch = cute.make_rmem_tensor((2,), cutlass.Int32) + cute.recast_tensor(scratch, cutlass.Int64)[0] = packed + return cls(flat_topk_index=scratch[0], topk_score=cute.recast_tensor(scratch, cutlass.Float32)[1]) + + +@cute.jit +def _copy_atom(dtype, num_bits_per_copy: int): + return cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), dtype, num_bits_per_copy=num_bits_per_copy) + + +@cute.jit +def _mark_alignment(tensor: cute.Tensor, byte_alignment: int) -> cute.Tensor: + pointer = tensor.iterator + return cute.make_tensor( + cute.make_ptr(pointer.dtype, pointer.toint(), pointer.memspace, assumed_align=byte_alignment), tensor.layout + ) + + +@cute.jit +def _device_trap() -> None: + """Terminate the current kernel after its failure state is published.""" + llvm.inline_asm( + res=None, + operands_=[], + asm_string="trap;", + constraints="", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +class _MetadataPushRouter(KernelComponent): + """Sort and push routing metadata into each destination rank's final pool.""" + + router_smem_limit_bytes: ClassVar[int] = 227 * 1024 + router_warps_per_cta: ClassVar[int] = 16 + + sizes_by_rank_region = "nvlink.token_comm.sizes_by_rank" + sizes_region = "nvlink.token_comm.sizes" + sizes_ready_region = "nvlink.token_comm.sizes_ready" + metadata_ready_region = "nvlink.token_comm.metadata_ready" + sorted_metadata_region = "nvlink.token_comm.sorted_metadata" + sorted_scores_region = "nvlink.token_comm.sorted_scores" + pool_expert_base_region = "nvlink.token_comm.pool_expert_base" + token_src_metadata_region = "nvlink.token_comm.token_src_metadata" # nosec B105 + fc1_topk_scores_region = "nvlink.token_comm.fc1_topk_scores" + source_expert_base_region = "nvlink.token_comm.source_expert_base" + push_destination_base_region = "nvlink.token_comm.push_destination_base" + sorted_metadata_ready_region = "nvlink.token_comm.sorted_metadata_ready" + # Private continuation: zero means this (source token, topk slot) is absent. + input_dedup_routes_region = "nvlink.token_comm.input_dedup_routes" + push_table_ready_region = "nvlink.token_comm.push_table_ready" + router_histogram_done_region = "nvlink.token_comm.router_histogram_done" + router_cta_histograms_region = "nvlink.token_comm.router_cta_histograms" + source_base_ready_region = "nvlink.token_comm.source_base_ready" + + router_data_histogram_region = "nvlink.token_comm.router_smem.data_histogram" + router_data_totals_region = "nvlink.token_comm.router_smem.data_totals" + router_data_prefix_region = "nvlink.token_comm.router_smem.data_prefix" + router_data_warp_totals_region = "nvlink.token_comm.router_smem.data_warp_totals" + router_data_sorted_region = "nvlink.token_comm.router_smem.data_sorted" + router_data_base_region = "nvlink.token_comm.router_smem.data_base" + router_helper_size_matrix_region = "nvlink.token_comm.router_smem.helper_size_matrix" + router_helper_totals_region = "nvlink.token_comm.router_smem.helper_totals" + router_helper_prefix_region = "nvlink.token_comm.router_smem.helper_prefix" + router_helper_warp_totals_region = "nvlink.token_comm.router_smem.helper_warp_totals" + router_helper_load_mbarrier_region = "nvlink.token_comm.router_smem.helper_load_mbarrier" + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return { + "world_size": int, + "expert_count": int, + "topk": int, + "max_tokens_per_rank": int, + "max_recv_size_per_rank": int, + "apply_topk_at_fc1": bool, + } + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return {"token_padding_block": int, "promised_launchable_sm_count": int, "drop_on_overflow": bool, + "input_dedup_mode": OptionalRequirement(str), + "interleave_source_ranks": OptionalRequirement(bool), + "sf_padding_block": OptionalRequirement(int), + "tokens_per_fc1_ready_slot": OptionalRequirement(int)} + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + self.world_size = problem_desc["world_size"] + self.expert_count = problem_desc["expert_count"] + self.topk = problem_desc["topk"] + self.max_tokens_per_rank = problem_desc["max_tokens_per_rank"] + self.max_recv_size_per_rank = min( + problem_desc["max_recv_size_per_rank"], self.world_size * self.max_tokens_per_rank * self.topk + ) + self.apply_topk_at_fc1 = problem_desc["apply_topk_at_fc1"] + + self.token_padding_block = impl_desc["token_padding_block"] + self.promised_launchable_sm_count = impl_desc["promised_launchable_sm_count"] + self.drop_on_overflow = impl_desc["drop_on_overflow"] + self.interleave_source_ranks = impl_desc.get("interleave_source_ranks", False) + + self.input_dedup_mode = impl_desc.get("input_dedup_mode", "off") + if self.input_dedup_mode not in ("off", "rank"): + raise ValueError("input_dedup_mode must be off/rank") + # The retained specialization uses expert-major ownership and publishes + # readiness only after every local fanout store has completed. + self.input_dedup_store_mode = "batched" if self.input_dedup_mode != "off" else "serial" + self.input_dedup_order = "expert" if self.input_dedup_mode != "off" else "token" + self._validate_router_configuration() + token_capacity = self.receive_capacity(self.token_padding_block) + if self.input_dedup_mode != "off": + # Route encoding uses three 20-bit row/ready fields. The smallest + # destination row selects the owner without encoding an expert ID. + if not self.apply_topk_at_fc1: + raise ValueError("Input dedup requires backward routing") + self.sf_padding_block = impl_desc["sf_padding_block"] + self.tokens_per_fc1_ready_slot = impl_desc["tokens_per_fc1_ready_slot"] + if self.sf_padding_block <= 0 or self.tokens_per_fc1_ready_slot <= 0: + raise ValueError("Input dedup requires positive SF padding and ready-slot granularity") + if self.max_recv_size_per_rank != self.world_size * self.max_tokens_per_rank * self.topk: + raise ValueError("Input dedup requires uncapped receive capacity") + # The low field stores pool_row + 1; zero denotes an unused route. + if token_capacity.padded_route_count >= (1 << 20): + raise ValueError("Input dedup requires padded receive capacity < 2**20 rows") + # SF rows and ready slots have separate zero-based 20-bit fields. + sf_capacity = self.receive_capacity(self.sf_padding_block).padded_route_count + ready_slots = ( + self.receive_capacity(self.tokens_per_fc1_ready_slot).padded_route_count + // self.tokens_per_fc1_ready_slot + ) + if sf_capacity > (1 << 20): + raise ValueError("Input dedup requires padded SF capacity <= 2**20 rows") + if ready_slots > (1 << 20): + raise ValueError("Input dedup requires at most 2**20 ready slots") + self.raw_route_count = token_capacity.raw_route_count + self.logical_route_capacity = token_capacity.logical_route_count + self.worst_case_token_count = token_capacity.padded_route_count + self.expert_count_padded = round_up(self.expert_count, 4) + self.expert_count_with_trash = self.expert_count_padded + 1 + self.router_elements_per_lane, self.router_data_cta_count = self._router_launch_configuration() + self.router_tokens_per_cta = self.router_elements_per_lane * self.router_warps_per_cta * 32 + self.router_push_cta_count = ceil_div(self.expert_count, self.router_warps_per_cta) + self.router_grid_cta_count = max(self.router_data_cta_count + 1, self.router_push_cta_count) + if self.router_grid_cta_count > self.promised_launchable_sm_count: + raise ValueError( + "Router grid exceeds promised_launchable_sm_count; all metadata-push CTAs must be concurrently resident." + ) + self._router_smem_workspace = self._build_router_smem_workspace() + + self._device_workspace = None + self._peer_rank_ptr_mapper = None + self._router_local_rank = None + self._router_thread_idx = None + self._router_linear_cta_idx = None + self._router_grid_thread_idx = None + self._router_warp_idx = None + self._router_lane_idx = None + self._overflow_flag = None + + def _validate_router_configuration(self) -> None: + if type(self.max_recv_size_per_rank) is not int: + raise TypeError(f"max_recv_size_per_rank must be an int, got {type(self.max_recv_size_per_rank).__name__}.") + if type(self.drop_on_overflow) is not bool: + raise TypeError(f"drop_on_overflow must be a bool, got {type(self.drop_on_overflow).__name__}.") + positive_fields = ( + "world_size", + "expert_count", + "topk", + "max_tokens_per_rank", + "max_recv_size_per_rank", + "token_padding_block", + "promised_launchable_sm_count", + ) + for field_name in positive_fields: + value = getattr(self, field_name) + if value <= 0: + raise ValueError(f"{field_name} must be positive, got {value}.") + if self.expert_count % self.world_size != 0: + raise ValueError( + f"expert_count must be divisible by world_size, got {self.expert_count} and {self.world_size}." + ) + if self.expert_count > 16384: + raise NotImplementedError("TokenComm supports at most 16384 global experts.") + if self.topk > self.expert_count: + raise ValueError(f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}.") + + @property + def experts_per_rank(self) -> int: + return self.expert_count // self.world_size + + def receive_capacity(self, padding_block: int) -> _ReceiveCapacity: + return _compute_receive_capacity( + world_size=self.world_size, + max_tokens_per_rank=self.max_tokens_per_rank, + topk=self.topk, + experts_per_rank=self.experts_per_rank, + max_recv_size_per_rank=self.max_recv_size_per_rank, + padding_block=padding_block, + ) + + def _router_launch_configuration(self) -> Tuple[int, int]: + def next_power_of_two(value: int) -> int: + return 1 << (max(value, 1) - 1).bit_length() + + routed_token_capacity = self.max_tokens_per_rank * self.topk + minimum_cta_capacity = 2048 + maximum_cta_capacity = 16384 + maximum_data_cta_count = 128 + maximum_supported_tokens = maximum_cta_capacity * maximum_data_cta_count + if routed_token_capacity > maximum_supported_tokens: + raise NotImplementedError(f"The router supports at most {maximum_supported_tokens} routed tokens per rank.") + cta_capacity = min(maximum_cta_capacity, next_power_of_two(max(routed_token_capacity, minimum_cta_capacity))) + elements_per_lane = cta_capacity // (self.router_warps_per_cta * 32) + data_cta_count = ceil_div(routed_token_capacity, cta_capacity) + return elements_per_lane, data_cta_count + + def _build_router_smem_workspace(self) -> SmemWorkspace: + workspace = SmemWorkspace() + workspace.register_mbarrier(self.router_helper_load_mbarrier_region, 1) + overlay = workspace.create_overlay("nvlink.token_comm.router_smem.role") + data_lifetime = overlay.add_lifetime("data_cta") + data_lifetime.register_tensor( + self.router_data_histogram_region, + cutlass.Int32, + (self.expert_count_with_trash, self.router_warps_per_cta + 1), + stride=(self.router_warps_per_cta + 1, 1), + ) + data_lifetime.register_tensor( + self.router_data_totals_region, cutlass.Int32, (self.expert_count_with_trash,), byte_alignment=16 + ) + data_lifetime.register_tensor( + self.router_data_prefix_region, cutlass.Int32, (self.expert_count_with_trash,), byte_alignment=16 + ) + data_lifetime.register_tensor(self.router_data_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,)) + data_lifetime.register_tensor( + self.router_data_sorted_region, + (cutlass.Int64 if self.apply_topk_at_fc1 else cutlass.Int32), + (self.router_tokens_per_cta,), + byte_alignment=16, + ) + if self.router_data_cta_count > 1: + data_lifetime.register_tensor( + self.router_data_base_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + + helper_lifetime = overlay.add_lifetime("helper_cta") + helper_lifetime.register_tensor( + self.router_helper_size_matrix_region, + cutlass.Int32, + (self.world_size, self.expert_count_padded), + stride=(self.expert_count_padded, 1), + byte_alignment=16, + ) + helper_lifetime.register_tensor( + self.router_helper_totals_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + helper_lifetime.register_tensor( + self.router_helper_prefix_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + helper_lifetime.register_tensor( + self.router_helper_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,), byte_alignment=16 + ) + workspace.finalize(max_bytes=self.router_smem_limit_bytes) + return workspace + + @property + def router_smem_workspace(self) -> SmemWorkspace: + return self._router_smem_workspace + + def register_device_workspace(self, workspace: DeviceWorkspace) -> None: + """Register router-private state and Router-to-Main outputs.""" + self._register_router_workspace(workspace) + + def _register_router_workspace(self, workspace: DeviceWorkspace) -> None: + maximum_routed_tokens = self.max_tokens_per_rank * self.topk + if self.input_dedup_mode != "off": + workspace.register( + self.input_dedup_routes_region, cutlass.Int64, + (self.world_size * maximum_routed_tokens,), + buffer_space="shared", byte_alignment=128, reset="tail_reset", + ) + workspace.register( + self.sizes_by_rank_region, + cutlass.Int32, + (self.world_size, self.expert_count_padded), + buffer_space="shared", + stride=(self.expert_count_padded, 1), + ) + workspace.register( + self.sizes_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="shared", reset="tail_reset" + ) + workspace.register(self.sizes_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.metadata_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.sorted_metadata_region, cutlass.Int64, (maximum_routed_tokens,), buffer_space="local") + if self.apply_topk_at_fc1: + workspace.register( + self.sorted_scores_region, cutlass.Float32, (maximum_routed_tokens,), buffer_space="local" + ) + workspace.register( + self.token_src_metadata_region, + cutlass.Int64, + (self.worst_case_token_count,), + buffer_space="shared", + byte_alignment=16, + ) + if self.apply_topk_at_fc1: + workspace.register( + self.fc1_topk_scores_region, cutlass.Float32, (self.worst_case_token_count,), buffer_space="shared" + ) + workspace.register(self.pool_expert_base_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local") + workspace.register( + self.source_expert_base_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="local" + ) + workspace.register( + self.push_destination_base_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="local" + ) + workspace.register( + self.sorted_metadata_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + workspace.register(self.push_table_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset") + if self.router_data_cta_count > 1: + workspace.register( + self.router_cta_histograms_region, + cutlass.Int32, + (self.router_data_cta_count, self.expert_count_padded), + buffer_space="local", + stride=(self.expert_count_padded, 1), + ) + workspace.register( + self.router_histogram_done_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + workspace.register( + self.source_base_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "_MetadataPushRouter": + if values: + raise ValueError("_MetadataPushRouter carries no MLIR values.") + return self + + @cute.jit + def launch_router( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + overflow_flag: cute.Tensor, + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper_host, + device_workspace: DeviceWorkspace, + stream: cuda.CUstream, + ) -> None: + """Launch counting-sort DATA, size-exchange HELPER, and metadata PUSH roles.""" + if cutlass.const_expr(self.apply_topk_at_fc1 and topk_scores is None): + raise ValueError("apply_topk_at_fc1 requires router topk_scores.") + if cutlass.const_expr(overflow_flag.iterator.dtype is not cutlass.Int32): + raise TypeError("overflow_flag must use Int32 elements.") + if cutlass.const_expr(cute.size(overflow_flag) != 1): + raise ValueError("overflow_flag must contain exactly one element.") + peer_rank_ptr_mapper = peer_rank_ptr_mapper_host.make_device_object() + self._router_kernel( + topk_indices, + topk_scores, + overflow_flag, + local_rank, + local_workspace, + shared_workspace, + peer_rank_ptr_mapper, + device_workspace, + ).launch( + grid=[self.router_grid_cta_count, 1, 1], + block=[self.router_warps_per_cta * 32, 1, 1], + min_blocks_per_mp=1, + stream=stream, + ) + + @cute.kernel + def _router_kernel( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + overflow_flag: cute.Tensor, + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper: SymmetricBufferDevice, + device_workspace: DeviceWorkspace, + ) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + linear_cta_idx, _, _ = cute.arch.block_idx() + cute.arch.griddepcontrol_launch_dependents() + block_thread_count = self.router_warps_per_cta * 32 + grid_thread_idx = thread_idx + linear_cta_idx * block_thread_count + warp_idx = cute.arch.make_warp_uniform(thread_idx // Int32(32)) + lane_idx = thread_idx % Int32(32) + + storage_type = self._router_smem_workspace.storage_class() + smem_allocator = cutlass.utils.SmemAllocator() + storage = smem_allocator.allocate(storage_type) + smem_base = storage.buffer.data_ptr() + + device_workspace.assign_device_members(local_workspace, shared_workspace) + self._device_workspace = device_workspace + self._peer_rank_ptr_mapper = peer_rank_ptr_mapper + self._router_local_rank = local_rank + self._router_thread_idx = thread_idx + self._router_linear_cta_idx = linear_cta_idx + self._router_grid_thread_idx = grid_thread_idx + self._router_warp_idx = warp_idx + self._router_lane_idx = lane_idx + self._overflow_flag = overflow_flag + + if cutlass.const_expr(self.router_data_cta_count == 1): + self._router_single_cta(topk_indices, topk_scores, smem_base) + else: + self._router_multiple_ctas(topk_indices, topk_scores, smem_base) + if linear_cta_idx < Int32(self.router_push_cta_count): + self._router_push_metadata() + + device_workspace.remove_device_members() + self._device_workspace = None + self._peer_rank_ptr_mapper = None + self._router_local_rank = None + self._router_thread_idx = None + self._router_linear_cta_idx = None + self._router_grid_thread_idx = None + self._router_warp_idx = None + self._router_lane_idx = None + self._overflow_flag = None + + @cute.jit + def _router_single_cta( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor], smem_base: cute.Pointer + ) -> None: + if self._router_linear_cta_idx < Int32(self.router_data_cta_count): + block_thread_count = self.router_warps_per_cta * 32 + trash_bucket = self.expert_count_padded + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) + totals = self._router_smem_workspace.tensor(self.router_data_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) + + histogram_element_count = self.expert_count_with_trash * (self.router_warps_per_cta + 1) + histogram_flat = cute.make_tensor(histogram.iterator, cute.make_layout((histogram_element_count,))) + zero_round_count = ceil_div(histogram_element_count, block_thread_count) + for zero_round in cutlass.range_constexpr(zero_round_count): + slot = Int32(zero_round * block_thread_count) + self._router_thread_idx + if slot < Int32(histogram_element_count): + histogram_flat[slot] = Int32(0) + + iket.range_push("router.histogram") + expert_registers, score_registers = self._load_router_inputs(topk_indices, topk_scores) + cute.arch.sync_threads() + match_masks = self._build_histogram(expert_registers, histogram) + self._prefix_warp_histogram(histogram, totals) + iket.range_pop() + + iket.range_push("router.prefix_and_publish") + publish_sizes = self._broadcast_sizes_to_peers( + cute.make_tensor(totals.iterator, cute.make_layout((self.expert_count_padded,))) + ) + total_valid_routes = smem_exclusive_prefix( + cute.make_tensor(totals.iterator, cute.make_layout((self.expert_count_padded,))), + cute.make_tensor(prefix.iterator, cute.make_layout((self.expert_count_padded,))), + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + if self._router_thread_idx == Int32(0): + prefix[trash_bucket] = total_valid_routes + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + source_expert_base[expert] = prefix[expert] + cute.arch.sync_threads() + publish_sizes() + iket.range_pop() + + iket.range_push("router.sort") + self._sort_router_elements( + expert_registers, match_masks, score_registers, sorted_elements, prefix, histogram, topk_indices.dtype + ) + cute.arch.sync_threads() + iket.range_pop() + iket.range_push("router.write_out") + self._dump_contiguous_router_output(sorted_elements, total_valid_routes) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" + ) + iket.range_pop() + elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): + self._router_helper_single_cta(smem_base) + + @cute.jit + def _router_multiple_ctas( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor], smem_base: cute.Pointer + ) -> None: + if self._router_linear_cta_idx < Int32(self.router_data_cta_count): + block_thread_count = self.router_warps_per_cta * 32 + trash_bucket = self.expert_count_padded + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) + totals = self._router_smem_workspace.tensor(self.router_data_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) + dump_base = self._router_smem_workspace.tensor(self.router_data_base_region, smem_base) + + histogram_element_count = self.expert_count_with_trash * (self.router_warps_per_cta + 1) + histogram_flat = cute.make_tensor(histogram.iterator, cute.make_layout((histogram_element_count,))) + zero_round_count = ceil_div(histogram_element_count, block_thread_count) + for zero_round in cutlass.range_constexpr(zero_round_count): + slot = Int32(zero_round * block_thread_count) + self._router_thread_idx + if slot < Int32(histogram_element_count): + histogram_flat[slot] = Int32(0) + + iket.range_push("router.histogram") + expert_registers, score_registers = self._load_router_inputs(topk_indices, topk_scores) + cute.arch.sync_threads() + match_masks = self._build_histogram(expert_registers, histogram) + self._prefix_warp_histogram(histogram, totals) + iket.range_pop() + + iket.range_push("router.reserve_and_prefix") + cta_histograms = self._device_workspace.tensor(self.router_cta_histograms_region) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + cta_histograms[self._router_linear_cta_idx, expert] = totals[expert] + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.router_histogram_done_region), Int32(1), sem="release", scope="gpu" + ) + + total_valid_routes = smem_exclusive_prefix( + cute.make_tensor(totals.iterator, cute.make_layout((self.expert_count_padded,))), + cute.make_tensor(prefix.iterator, cute.make_layout((self.expert_count_padded,))), + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + if self._router_thread_idx == Int32(0): + prefix[trash_bucket] = total_valid_routes + cute.arch.sync_threads() + iket.range_pop() + iket.range_push("router.sort") + self._sort_router_elements( + expert_registers, match_masks, score_registers, sorted_elements, prefix, histogram, topk_indices.dtype + ) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.wait_source_base") + source_base_ready = self._device_workspace.ptr(self.source_base_ready_region) + if self._router_thread_idx == Int32(0): + while cute.arch.load(source_base_ready, Int32, sem="acquire", scope="gpu") != Int32(1): + nanosleep(150) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.write_out") + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + cta_histograms = self._device_workspace.tensor(self.router_cta_histograms_region) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + dump_base[expert] = source_expert_base[expert] + cta_histograms[self._router_linear_cta_idx, expert] + cute.arch.sync_threads() + self._dump_router_output_by_expert(totals, prefix, dump_base, sorted_elements) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" + ) + iket.range_pop() + elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): + self._router_helper_multiple_ctas(smem_base) + + @cute.jit + def _router_helper_single_cta(self, smem_base: cute.Pointer) -> None: + iket.range_push("router.compute_push_tables") + self._compute_push_tables(smem_base) + iket.range_pop() + + @cute.jit + def _router_helper_multiple_ctas(self, smem_base: cute.Pointer) -> None: + block_thread_count = self.router_warps_per_cta * 32 + totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) + cta_histograms = self._device_workspace.tensor(self.router_cta_histograms_region) + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + + histogram_done = self._device_workspace.ptr(self.router_histogram_done_region) + iket.range_push("router.wait_histogram") + if self._router_thread_idx == Int32(0): + while cute.arch.load(histogram_done, Int32, sem="acquire", scope="gpu") != Int32( + self.router_data_cta_count + ): + nanosleep(150) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.broadcast_sizes") + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + running = Int32(0) + for cta_idx in cutlass.range_constexpr(self.router_data_cta_count): + count = cta_histograms[cta_idx, expert] + cta_histograms[cta_idx, expert] = running + running = running + count + totals[expert] = running + cute.arch.sync_threads() + + publish_sizes = self._broadcast_sizes_to_peers(totals) + smem_exclusive_prefix( + totals, + prefix, + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + source_expert_base[expert] = prefix[expert] + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.source_base_ready_region), Int32(1), sem="release", scope="gpu" + ) + publish_sizes() + iket.range_pop() + + iket.range_push("router.compute_push_tables") + self._compute_push_tables(smem_base) + iket.range_pop() + + @cute.jit + def _router_push_metadata(self) -> None: + block_thread_count = self.router_warps_per_cta * 32 + sorted_metadata_ready = self._device_workspace.ptr(self.sorted_metadata_ready_region) + push_table_ready = self._device_workspace.ptr(self.push_table_ready_region) + if self._router_thread_idx == Int32(0): + while cute.arch.load(sorted_metadata_ready, Int32, sem="acquire", scope="gpu") != Int32( + self.router_data_cta_count + ): + nanosleep(150) + while cute.arch.load(push_table_ready, Int32, sem="acquire", scope="gpu") != Int32(1): + nanosleep(150) + cute.arch.sync_threads() + + global_expert = self._router_linear_cta_idx * Int32(self.router_warps_per_cta) + self._router_warp_idx + if global_expert < Int32(self.expert_count): + sizes_by_rank = self._device_workspace.tensor(self.sizes_by_rank_region) + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + push_destination_base = self._device_workspace.tensor(self.push_destination_base_region) + route_count = sizes_by_rank[self._router_local_rank, global_expert] + source_begin = source_expert_base[global_expert] + destination_begin = push_destination_base[global_expert] + destination_rank = global_expert // Int32(self.experts_per_rank) + if cutlass.const_expr(self.interleave_source_ranks): + source_counts = [sizes_by_rank[r, global_expert] for r in range(self.world_size)] + peer_offset = self._peer_rank_ptr_mapper.map(Int64(0), destination_rank, Int64(0)) + source_metadata = self._device_workspace.ptr(self.sorted_metadata_region) + destination_metadata_address = ( + self._device_workspace.ptr(self.token_src_metadata_region).toint() + peer_offset + ) + if cutlass.const_expr(self.apply_topk_at_fc1): + source_scores = self._device_workspace.ptr(self.sorted_scores_region) + destination_scores_address = ( + self._device_workspace.ptr(self.fc1_topk_scores_region).toint() + peer_offset + ) + if cutlass.const_expr(self.input_dedup_mode != "off"): + destination_routes_address = self._device_workspace.ptr(self.input_dedup_routes_region).toint() + peer_offset + local_expert = global_expert % Int32(self.experts_per_rank) + global_sizes = self._device_workspace.tensor(self.sizes_region) + pool_begin = Int32(0) + sf_begin = Int32(0) + ready_begin = Int32(0) + for earlier in cutlass.range_constexpr(self.experts_per_rank): + if Int32(earlier) < local_expert: + count = global_sizes[destination_rank * Int32(self.experts_per_rank) + Int32(earlier)] + pool_begin = pool_begin + ( + (count + Int32(self.token_padding_block - 1)) // Int32(self.token_padding_block) + ) * Int32(self.token_padding_block) + sf_begin = sf_begin + ( + (count + Int32(self.sf_padding_block - 1)) // Int32(self.sf_padding_block) + ) * Int32(self.sf_padding_block) + ready_begin = ready_begin + ( + (count + Int32(self.tokens_per_fc1_ready_slot - 1)) + // Int32(self.tokens_per_fc1_ready_slot) + ) + route_round_count = (route_count + Int32(31)) // Int32(32) + for route_round in cutlass.range(route_round_count, unroll=1): + route = Int32(route_round) * Int32(32) + self._router_lane_idx + if route < route_count: + source_position = source_begin + route + destination_position = destination_begin + route + if cutlass.const_expr(self.interleave_source_ranks): + # Visit source ranks in order, skipping exhausted ranks. + # This is also the canonical row order of fc1_preact: + # FWD and BWD must select the same ordering. + round_robin_offset = Int32(0) + for source_rank in cutlass.range_constexpr(self.world_size): + count = source_counts[source_rank] + round_robin_offset = round_robin_offset + cutlass.min(route, count) + if Int32(source_rank) < Int32(self._router_local_rank) and route < count: + round_robin_offset = round_robin_offset + Int32(1) + destination_position = destination_begin + round_robin_offset + physical_store = Int32(destination_position < Int32(self.worst_case_token_count)) + metadata = cute.arch.load(source_metadata + source_position, cutlass.Int64) + stg_b64( + destination_metadata_address + Int64(destination_position) * Int64(TokenSrcMetadata.nbytes), + metadata, + physical_store, + ) + if cutlass.const_expr(self.input_dedup_mode != "off"): + source_token = Int32(metadata & Int64(0xFFFFFFFF)) + source_slot = Int32((metadata >> Int64(32)) & Int64(0xFFFF)) + route_index = (self._router_local_rank * Int32(self.max_tokens_per_rank) + source_token) * Int32(self.topk) + source_slot + token_in_expert = destination_position - pool_begin + sf_position = sf_begin + token_in_expert + ready_slot = ready_begin + token_in_expert // Int32(self.tokens_per_fc1_ready_slot) + # Zero denotes an absent route; the destination row + # also gives each source token/receiver a unique owner. + packed_route = ( + (Int64(ready_slot) << Int64(40)) | (Int64(sf_position) << Int64(20)) + | Int64(destination_position + Int32(1)) + ) + stg_b64(destination_routes_address + Int64(route_index) * Int64(8), packed_route, physical_store) + if cutlass.const_expr(self.apply_topk_at_fc1): + score = cute.arch.load(source_scores + source_position, cutlass.Float32) + stg_f32( + destination_scores_address + Int64(destination_position) * Int64(4), score, physical_store + ) + + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_sys() + # Keep notifier threads behind the leader's system fence. + cute.arch.sync_threads() + metadata_ready_address = self._device_workspace.ptr(self.metadata_ready_region).toint() + rank_round_count = ceil_div(self.world_size, block_thread_count) + for rank_round in cutlass.range_constexpr(rank_round_count): + destination_rank = Int32(rank_round * block_thread_count) + self._router_thread_idx + if destination_rank < Int32(self.world_size): + red_add_relaxed_sys_s32( + self._peer_rank_ptr_mapper.map(metadata_ready_address, destination_rank, Int64(0)), Int32(1) + ) + + @cute.jit + def _broadcast_sizes_to_peers(self, smem_expert_counts: cute.Tensor) -> Callable[[], None]: + block_thread_count = self.router_warps_per_cta * 32 + row_bytes = Int32(self.expert_count_padded * 4) + matrix_address = self._device_workspace.ptr(self.sizes_by_rank_region).toint() + total_address = self._device_workspace.ptr(self.sizes_region).toint() + rank_round_count = ceil_div(self.world_size, self.router_warps_per_cta) + for rank_round in cutlass.range_constexpr(rank_round_count): + destination_rank = self._router_warp_idx + Int32(rank_round * self.router_warps_per_cta) + if destination_rank < Int32(self.world_size): + peer_offset = self._peer_rank_ptr_mapper.map(Int64(0), destination_rank, Int64(0)) + destination_row_address = ( + matrix_address + + peer_offset + + Int64(Int32(self._router_local_rank) * Int32(self.expert_count_padded)) * Int64(4) + ) + destination_row = cute.make_ptr( + cutlass.Int32, destination_row_address, AddressSpace.gmem, assumed_align=16 + ) + destination_total = cute.make_ptr( + cutlass.Int32, total_address + peer_offset, AddressSpace.gmem, assumed_align=16 + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_row, smem_expert_counts.iterator, row_bytes) + cp_reduce_async_bulk_add_u32_s2g(destination_total, smem_expert_counts.iterator, row_bytes) + cute.arch.cp_async_bulk_commit_group() + + def finalize() -> None: + cute.arch.cp_async_bulk_wait_group(0) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_sys() + cute.arch.sync_threads() + ready_address = self._device_workspace.ptr(self.sizes_ready_region).toint() + ready_round_count = ceil_div(self.world_size, block_thread_count) + for ready_round in cutlass.range_constexpr(ready_round_count): + destination_rank = Int32(ready_round * block_thread_count) + self._router_thread_idx + if destination_rank < Int32(self.world_size): + red_add_relaxed_sys_s32( + self._peer_rank_ptr_mapper.map(ready_address, destination_rank, Int64(0)), Int32(1) + ) + + return finalize + + @cute.jit + def _compute_push_tables(self, smem_base: cute.Pointer) -> None: + block_thread_count = self.router_warps_per_cta * 32 + owner_expert_begin = Int32(self._router_local_rank) * Int32(self.experts_per_rank) + matrix_bytes = self.world_size * self.expert_count_padded * 4 + + size_matrix = self._router_smem_workspace.tensor(self.router_helper_size_matrix_region, smem_base) + padded_totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) + load_mbarrier = self._router_smem_workspace.ptr(self.router_helper_load_mbarrier_region, smem_base) + sizes = self._device_workspace.tensor(self.sizes_region) + + if self._router_thread_idx == Int32(0): + cute.arch.mbarrier_init(load_mbarrier, 1) + + sizes_ready = self._device_workspace.ptr(self.sizes_ready_region) + iket.range_push("router.wait_sizes_ready") + if self._router_thread_idx == Int32(0): + while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="sys") != Int32(self.raw_sizes_ready_target): + nanosleep(150) + cute.arch.mbarrier_init_fence() + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.load_sizes_and_prefix") + if self._router_thread_idx == Int32(0): + cute.arch.mbarrier_arrive_and_expect_tx(load_mbarrier, Int32(matrix_bytes)) + tma_load_1d( + size_matrix.iterator, + self._device_workspace.ptr(self.sizes_by_rank_region), + load_mbarrier, + Int32(matrix_bytes), + ) + + padded_expert_rounds = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(padded_expert_rounds): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + expert_size = sizes[expert] + padded_totals[expert] = ( + (expert_size + Int32(self.token_padding_block - 1)) // Int32(self.token_padding_block) + ) * Int32(self.token_padding_block) + cute.arch.sync_threads() + smem_exclusive_prefix( + padded_totals, + prefix, + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + pool_expert_base = self._device_workspace.tensor(self.pool_expert_base_region) + local_expert_rounds = ceil_div(self.experts_per_rank, block_thread_count) + for expert_round in cutlass.range_constexpr(local_expert_rounds): + local_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if local_expert < Int32(self.experts_per_rank): + pool_expert_base[local_expert] = prefix[owner_expert_begin + local_expert] - prefix[owner_expert_begin] + + cute.arch.mbarrier_wait(load_mbarrier, 0) + iket.range_pop() + + iket.range_push("router.build_push_destinations") + push_destination_base = self._device_workspace.tensor(self.push_destination_base_region) + padded_expert_rounds = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(padded_expert_rounds): + global_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if global_expert < Int32(self.expert_count): + destination_rank = global_expert // Int32(self.experts_per_rank) + destination_expert_begin = destination_rank * Int32(self.experts_per_rank) + destination_pool_base = prefix[global_expert] - prefix[destination_expert_begin] + source_ring_offset = Int32(0) + if cutlass.const_expr(not self.interleave_source_ranks): + local_ring_position = ( + Int32(self._router_local_rank) - destination_rank + Int32(self.world_size) + ) % Int32(self.world_size) + for ring_position in cutlass.range_constexpr(self.world_size): + source_rank = (destination_rank + Int32(ring_position)) % Int32(self.world_size) + if Int32(ring_position) < local_ring_position: + source_ring_offset = source_ring_offset + size_matrix[source_rank, global_expert] + push_destination_base[global_expert] = destination_pool_base + source_ring_offset + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.apply_receive_limit") + if self._router_thread_idx == Int32(0): + warp_totals[0] = Int32(0) + cute.arch.sync_threads() + thread_local_total = Int32(0) + local_expert_rounds = ceil_div(self.experts_per_rank, block_thread_count) + for expert_round in cutlass.range_constexpr(local_expert_rounds): + local_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if local_expert < Int32(self.experts_per_rank): + thread_local_total = thread_local_total + sizes[owner_expert_begin + local_expert] + if thread_local_total > Int32(0): + cute.arch.atomic_add(warp_totals.iterator, thread_local_total, scope="cta") + cute.arch.sync_threads() + raw_local_total = warp_totals[0] + did_overflow = Int32(raw_local_total > Int32(self.max_recv_size_per_rank)) + + if cutlass.const_expr(self.drop_on_overflow): + for expert_round in cutlass.range_constexpr(local_expert_rounds): + local_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if local_expert < Int32(self.experts_per_rank): + # Keep every raw row that still lands inside the statically allocated padded pool. + raw_size = sizes[owner_expert_begin + local_expert] + retained_size = Int32(0) + remaining_capacity = Int32(self.worst_case_token_count) - pool_expert_base[local_expert] + if remaining_capacity > Int32(0): + retained_size = cutlass.min(raw_size, remaining_capacity) + if did_overflow: + sizes[owner_expert_begin + local_expert] = retained_size + + if self._router_thread_idx == Int32(0): + cute.arch.store(self._overflow_flag.iterator, did_overflow, sem="relaxed", scope="sys") + cute.arch.sync_threads() + + if cutlass.const_expr(self.drop_on_overflow): + self._publish_push_table_and_sizes_ready(sizes_ready) + else: + if did_overflow: + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_sys() + _device_trap() + else: + self._publish_push_table_and_sizes_ready(sizes_ready) + iket.range_pop() + + @cute.jit + def _publish_push_table_and_sizes_ready(self, sizes_ready: cute.Pointer) -> None: + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_gpu() + cute.arch.atomic_add( + self._device_workspace.ptr(self.push_table_ready_region), Int32(1), sem="relaxed", scope="gpu" + ) + cute.arch.atomic_add(sizes_ready, Int32(1), sem="relaxed", scope="gpu") + + @cute.jit + def _load_router_inputs( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor] + ) -> Tuple[cute.Tensor, Optional[cute.Tensor]]: + elements_per_vector = 128 // topk_indices.dtype.width + grid_thread_count = self.router_data_cta_count * self.router_warps_per_cta * 32 + tile_span = elements_per_vector * grid_thread_count + maximum_elements = self.max_tokens_per_rank * self.topk + actual_token_count = Int32(self.max_tokens_per_rank) + actual_elements = Int32(maximum_elements) + load_round_count = ceil_div(maximum_elements, tile_span) + elements_per_thread = load_round_count * elements_per_vector + + topk_flat = cute.make_tensor(topk_indices.iterator, cute.make_layout((maximum_elements,))) + topk_vectors = cute.logical_divide(cute.zipped_divide(topk_flat, (tile_span,)), (elements_per_vector, None)) + load_atom = _copy_atom(topk_indices.dtype, 128) + expert_registers = cute.make_rmem_tensor((elements_per_thread,), cutlass.Int32) + if cutlass.const_expr(topk_indices.dtype.width == 64): + raw_indices = cute.make_rmem_tensor((elements_per_thread,), topk_indices.dtype) + raw_vectors = cute.zipped_divide(raw_indices, (elements_per_vector,)) + for load_round in cutlass.range_constexpr(load_round_count): + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) + if tile_begin < actual_elements: + cute.copy( + load_atom, + _mark_alignment(topk_vectors[(None, self._router_grid_thread_idx), load_round], 16), + raw_vectors[None, load_round], + ) + else: + expert_vectors = cute.zipped_divide(expert_registers, (elements_per_vector,)) + for load_round in cutlass.range_constexpr(load_round_count): + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) + if tile_begin < actual_elements: + cute.copy( + load_atom, + _mark_alignment(topk_vectors[(None, self._router_grid_thread_idx), load_round], 16), + expert_vectors[None, load_round], + ) + + score_registers = None + if cutlass.const_expr(self.apply_topk_at_fc1): + score_registers = cute.make_rmem_tensor((elements_per_thread,), cutlass.Float32) + scores_flat = cute.make_tensor(topk_scores.iterator, cute.make_layout((maximum_elements,))) + score_vectors = cute.logical_divide( + cute.zipped_divide(scores_flat, (tile_span,)), (elements_per_vector, None) + ) + score_atom = _copy_atom(cutlass.Float32, elements_per_vector * 32) + score_register_vectors = cute.zipped_divide(score_registers, (elements_per_vector,)) + for load_round in cutlass.range_constexpr(load_round_count): + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) + if tile_begin < actual_elements: + cute.copy( + score_atom, + _mark_alignment(score_vectors[(None, self._router_grid_thread_idx), load_round], 16), + score_register_vectors[None, load_round], + ) + + expert_registers_u32 = cute.recast_tensor(expert_registers, cutlass.Uint32) + if cutlass.const_expr(topk_indices.dtype.width == 64): + raw_indices_i32 = cute.recast_tensor(raw_indices, cutlass.Int32) + for register_idx in cutlass.range_constexpr(elements_per_thread): + if cutlass.const_expr(topk_indices.dtype.width == 64): + expert_registers[register_idx] = raw_indices_i32[2 * register_idx] + token_idx, _ = self._router_value_coordinate(register_idx, topk_indices.dtype) + is_invalid = (expert_registers_u32[register_idx] >= cutlass.Uint32(self.expert_count)) | ( + token_idx >= actual_token_count + ) + if is_invalid: + expert_registers[register_idx] = Int32(self.expert_count_padded) + return expert_registers, score_registers + + @cute.jit + def _build_histogram(self, expert_registers: cute.Tensor, histogram: cute.Tensor) -> cute.Tensor: + register_count = cute.size(expert_registers) + match_masks = cute.make_rmem_tensor((register_count,), cutlass.Int32) + lane_mask_less_than = (Int32(1) << self._router_lane_idx) - Int32(1) + for register_idx in cutlass.range_constexpr(register_count): + expert = expert_registers[register_idx] + match_mask = Int32(cute.arch.match_sync(0xFFFFFFFF, expert, kind="any")) + match_masks[register_idx] = match_mask + histogram_slot = expert * Int32(self.router_warps_per_cta + 1) + self._router_warp_idx + rank_in_group = Int32(cute.arch.popc(match_mask & lane_mask_less_than)) + if rank_in_group == Int32(0): + cute.arch.atomic_add( + histogram.iterator + histogram_slot, Int32(cute.arch.popc(match_mask)), sem="relaxed", scope="cta" + ) + cute.arch.sync_threads() + return match_masks + + @cute.jit + def _prefix_warp_histogram(self, histogram: cute.Tensor, totals: cute.Tensor) -> None: + """Turn per-warp counts into fixed warp bases and expert totals.""" + block_thread_count = self.router_warps_per_cta * 32 + expert_round_count = ceil_div(self.expert_count_with_trash, block_thread_count) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_with_trash): + running = Int32(0) + for warp in cutlass.range_constexpr(self.router_warps_per_cta): + count = histogram[expert, warp] + histogram[expert, warp] = running + running = running + count + totals[expert] = running + cute.arch.sync_threads() + + @cute.jit + def _sort_router_elements( + self, + expert_registers: cute.Tensor, + match_masks: cute.Tensor, + score_registers: Optional[cute.Tensor], + sorted_elements: cute.Tensor, + expert_run_starts: cute.Tensor, + warp_histogram: cute.Tensor, + topk_index_type: type, + ) -> None: + """Stable CTA-local counting-sort scatter. + + The cursor for an expert is private to one warp. ``match.any`` gives + every equal-expert lane a fixed rank, and only that group's first lane + advances the cursor. Therefore the output order is a pure function of + warp id, register round, and lane id rather than atomic arrival order. + """ + register_count = cute.size(expert_registers) + lane_mask_less_than = (Int32(1) << self._router_lane_idx) - Int32(1) + for register_idx in cutlass.range_constexpr(register_count): + expert = expert_registers[register_idx] + match_mask = match_masks[register_idx] + group_size = Int32(cute.arch.popc(match_mask)) + rank_in_group = Int32(cute.arch.popc(match_mask & lane_mask_less_than)) + warp_base = warp_histogram[expert, self._router_warp_idx] + token_idx, topk_slot = self._router_value_coordinate(register_idx, topk_index_type) + flat_topk_index = token_idx * Int32(self.topk) + topk_slot + destination = expert_run_starts[expert] + warp_base + rank_in_group + + # Every lane in an equal-expert group reads the old cursor before + # its leader advances it for the next register round. + cute.arch.sync_warp() + if rank_in_group == Int32(0): + warp_histogram[expert, self._router_warp_idx] = warp_base + group_size + if cutlass.const_expr(self.apply_topk_at_fc1): + sorted_elements[destination] = _SortedElement(flat_topk_index, score_registers[register_idx]).pack() + else: + sorted_elements[destination] = _SortedElement(flat_topk_index, None).pack() + cute.arch.sync_warp() + + @cute.jit + def _router_value_coordinate(self, register_idx: int, topk_index_type: type) -> Tuple[Int32, Int32]: + elements_per_vector = 128 // topk_index_type.width + tile_span = elements_per_vector * self.router_data_cta_count * self.router_warps_per_cta * 32 + flat_index = Int32( + register_idx // elements_per_vector * tile_span + register_idx % elements_per_vector + ) + self._router_grid_thread_idx * Int32(elements_per_vector) + return (flat_index // Int32(self.topk), flat_index % Int32(self.topk)) + + @cute.jit + def _dump_contiguous_router_output(self, sorted_elements: cute.Tensor, total_valid_routes: Int32) -> None: + block_thread_count = self.router_warps_per_cta * 32 + metadata_address = self._device_workspace.ptr(self.sorted_metadata_region).toint() + if cutlass.const_expr(self.apply_topk_at_fc1): + score_address = self._device_workspace.ptr(self.sorted_scores_region).toint() + dump_round_count = (total_valid_routes + Int32(block_thread_count - 1)) // Int32(block_thread_count) + for dump_round in cutlass.range(dump_round_count, unroll=4): + position = Int32(dump_round * block_thread_count) + self._router_thread_idx + predicate = Int32(position < total_valid_routes) + element = _SortedElement.from_packed(sorted_elements[position]) + metadata = TokenSrcMetadata( + src_rank=Int32(self._router_local_rank), + src_token=(element.flat_topk_index // Int32(self.topk)), + src_topk=(element.flat_topk_index % Int32(self.topk)), + ) + stg_b64(metadata_address + Int64(position) * Int64(TokenSrcMetadata.nbytes), metadata.pack(), predicate) + if cutlass.const_expr(self.apply_topk_at_fc1): + stg_f32(score_address + Int64(position) * Int64(4), element.topk_score, predicate) + + @cute.jit + def _dump_router_output_by_expert( + self, + expert_totals: cute.Tensor, + expert_run_starts: cute.Tensor, + expert_dump_bases: cute.Tensor, + sorted_elements: cute.Tensor, + ) -> None: + metadata_address = self._device_workspace.ptr(self.sorted_metadata_region).toint() + if cutlass.const_expr(self.apply_topk_at_fc1): + score_address = self._device_workspace.ptr(self.sorted_scores_region).toint() + expert_round_count = ceil_div(self.expert_count_padded, self.router_warps_per_cta) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = self._router_warp_idx + Int32(expert_round * self.router_warps_per_cta) + if expert < Int32(self.expert_count_padded): + run_begin = expert_run_starts[expert] + run_length = expert_totals[expert] + dump_begin = expert_dump_bases[expert] + route_round_count = (run_length + Int32(31)) // Int32(32) + for route_round in cutlass.range(route_round_count, unroll=1): + route = Int32(route_round) * Int32(32) + self._router_lane_idx + predicate = Int32(route < run_length) + element = _SortedElement.from_packed(sorted_elements[predicate * (run_begin + route)]) + output_position = dump_begin + route + metadata = TokenSrcMetadata( + src_rank=Int32(self._router_local_rank), + src_token=(element.flat_topk_index // Int32(self.topk)), + src_topk=(element.flat_topk_index % Int32(self.topk)), + ) + stg_b64( + metadata_address + Int64(output_position) * Int64(TokenSrcMetadata.nbytes), + metadata.pack(), + predicate, + ) + if cutlass.const_expr(self.apply_topk_at_fc1): + stg_f32(score_address + Int64(output_position) * Int64(4), element.topk_score, predicate) + + @cute.jit + def local_expert_sizes(self, device_workspace: DeviceWorkspace, local_rank: Int32) -> cute.Tensor: + """Return this rank's contiguous expert-size view.""" + sizes = device_workspace.tensor(self.sizes_region) + expert_begin = local_rank * Int32(self.experts_per_rank) + return cute.make_tensor(sizes.iterator + expert_begin, cute.make_layout((self.experts_per_rank,))) + + @property + def metadata_ready_target(self) -> int: + return self.router_push_cta_count * self.world_size + + @property + def raw_sizes_ready_target(self) -> int: + return self.world_size + + @property + def published_sizes_ready_target(self) -> int: + return self.world_size + 1 + + @cute.jit + def sizes_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.sizes_region) + + @cute.jit + def pool_expert_base_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.pool_expert_base_region) + + @cute.jit + def token_src_metadata_pointer(self, device_workspace: DeviceWorkspace) -> cute.Pointer: + return device_workspace.ptr(self.token_src_metadata_region) + + @cute.jit + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + if lane_idx == Int32(0): + sizes_ready = device_workspace.ptr(self.sizes_ready_region) + while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="gpu") != Int32( + self.published_sizes_ready_target + ): + nanosleep(sleep_cycles) + cute.arch.sync_warp() + + @cute.jit + def wait_for_metadata_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + if lane_idx == Int32(0): + metadata_ready = device_workspace.ptr(self.metadata_ready_region) + while cute.arch.load(metadata_ready, Int32, sem="acquire", scope="sys") != Int32( + self.metadata_ready_target + ): + nanosleep(sleep_cycles) + cute.arch.sync_warp() + + @cute.jit + def token_src_metadata_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.token_src_metadata_region) + + @cute.jit + def fc1_topk_scores_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.apply_topk_at_fc1): + return None + return device_workspace.tensor(self.fc1_topk_scores_region) + + +class TokenCommDeterministic(KernelComponent): + """Fused communication with a fixed logical token-position sequence.""" + + transfer_warp_count: ClassVar[int] = 4 + transfer_thread_count: ClassVar[int] = transfer_warp_count * 32 + standalone_chunk_bytes: ClassVar[int] = 2048 + minimum_pacing_window_cycles: ClassVar[int] = 512 + standalone_max_backoff_cycles: ClassVar[int] = 500 + adaptive_minimum_sleep_cycles: ClassVar[int] = 50 + transfer_lifetime_barrier_id: ClassVar[int] = 9 + grid_sync_barrier_id: ClassVar[int] = 10 + standalone_size_barrier_id: ClassVar[int] = 11 + token_in_size_barrier_id: ClassVar[int] = 12 + + fc1_ready_region = "nvlink.token_comm.fc1_ready" + fc1_activation_region = "nvlink.token_comm.fc1_activation" + fc1_activation_sf_region = "nvlink.token_comm.fc1_activation_sf" + fc2_done_region = "nvlink.token_comm.fc2_done" + fc2_activation_region = "nvlink.token_comm.fc2_activation" + fc2_activation_sf_region = "nvlink.token_comm.fc2_activation_sf" + pre_reduced_activation_region = "nvlink.token_comm.pre_reduced_activation" + pre_reduced_activation_sf_region = "nvlink.token_comm.pre_reduced_activation_sf" + token_back_schedule_region = "nvlink.token_comm.token_back_schedule" # nosec B105 + + token_in_mbarrier_region = "nvlink.token_comm.main_smem.token_in_mbarriers" # nosec B105 + token_back_mbarrier_region = "nvlink.token_comm.main_smem.token_back_mbarriers" # nosec B105 + expert_sizes_smem_region = "nvlink.token_comm.main_smem.expert_sizes" + token_in_activation_smem_region = "nvlink.token_comm.main_smem.token_in_activation" # nosec B105 + token_in_sf_smem_region = "nvlink.token_comm.main_smem.token_in_sf" # nosec B105 + token_back_activation_smem_region = "nvlink.token_comm.main_smem.token_back_activation" # nosec B105 + token_back_sf_smem_region = "nvlink.token_comm.main_smem.token_back_sf" # nosec B105 + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return { + "world_size": int, + "expert_count": int, + "topk": int, + "max_tokens_per_rank": int, + "max_recv_size_per_rank": int, + "hidden_size": int, + "quant_kind": str, + "combine_format": CombineFormat, + "apply_topk_at_fc1": bool, + } + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return { + "token_padding_block": int, + "sf_padding_block": int, + "tokens_per_fc1_ready_slot": int, + "fc2_done_signals_per_token_tile": int, + "promised_launchable_sm_count": int, + # Optional so existing component users keep the four-warp ABI. + # Embedding kernels must still launch exactly this many transfer + # warps and include them in their block-wide synchronization. + "token_in_transfer_warp_count": OptionalRequirement(int), + "token_in_flag_batch": int, + "token_in_async_flags": OptionalRequirement(bool), + "input_dedup_mode": OptionalRequirement(str), + "interleave_source_ranks": OptionalRequirement(bool), + "token_back_mode": str, + "token_back_schedule_mode": str, + "reduce_topk_in_kernel": bool, + "drop_on_overflow": bool, + } + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + self.world_size = problem_desc["world_size"] + self.expert_count = problem_desc["expert_count"] + self.topk = problem_desc["topk"] + self.max_tokens_per_rank = problem_desc["max_tokens_per_rank"] + self.max_recv_size_per_rank = min( + problem_desc["max_recv_size_per_rank"], self.world_size * self.max_tokens_per_rank * self.topk + ) + self.hidden_size = problem_desc["hidden_size"] + self.quant_kind = problem_desc["quant_kind"] + self.combine_format = problem_desc["combine_format"] + self.apply_topk_at_fc1 = problem_desc["apply_topk_at_fc1"] + + self.token_padding_block = impl_desc["token_padding_block"] + self.sf_padding_block = impl_desc["sf_padding_block"] + self.tokens_per_fc1_ready_slot = impl_desc["tokens_per_fc1_ready_slot"] + self.fc2_done_signals_per_token_tile = impl_desc["fc2_done_signals_per_token_tile"] + self.promised_launchable_sm_count = impl_desc["promised_launchable_sm_count"] + self.transfer_warp_count = impl_desc.get( + "token_in_transfer_warp_count", type(self).transfer_warp_count + ) + self.transfer_thread_count = self.transfer_warp_count * 32 + self.token_in_flag_batch = impl_desc["token_in_flag_batch"] + self.token_in_async_flags = impl_desc.get("token_in_async_flags", False) + self.token_back_mode: TokenBackMode = impl_desc["token_back_mode"] + self.token_back_schedule_mode: TokenBackScheduleMode = impl_desc["token_back_schedule_mode"] + self.reduce_topk_in_kernel = impl_desc["reduce_topk_in_kernel"] + self.drop_on_overflow = impl_desc["drop_on_overflow"] + + self._validate_configuration() + self._router = _MetadataPushRouter(problem_desc, impl_desc) + self.interleave_source_ranks = self._router.interleave_source_ranks + self.input_dedup_mode = self._router.input_dedup_mode + self.input_dedup_store_mode = self._router.input_dedup_store_mode + self.input_dedup_order = self._router.input_dedup_order + if self.input_dedup_mode != "off" and self.quant_kind not in ("mxfp8_e4m3", "mxfp8_e5m2"): + raise ValueError("Input dedup requires MXFP8") + self._nvlink_barrier = NvlinkBarrier(world_size=self.world_size, barrier_id=self.grid_sync_barrier_id) + self._device_workspace = None + self._token_comm_args = None + self._local_rank = None + self._linear_cta_idx = None + self._transfer_warp_idx = None + self._lane_idx = None + + def _validate_configuration(self) -> None: + positive_fields = ("hidden_size", "token_padding_block", "sf_padding_block", "tokens_per_fc1_ready_slot") + for field_name in positive_fields: + value = getattr(self, field_name) + if value <= 0: + raise ValueError(f"{field_name} must be positive, got {value}.") + if self.quant_kind not in _quant_spec: + raise ValueError(f"Unsupported quant_kind {self.quant_kind!r}.") + if self.token_back_mode not in ("epi_warps", "standalone_warps", "reuse_dispatch_warps"): + raise ValueError(f"Unsupported token_back_mode {self.token_back_mode!r}.") + if self.token_back_schedule_mode not in ("static", "atomic_counter"): + raise ValueError( + f"token_back_schedule_mode must be static or atomic_counter, got {self.token_back_schedule_mode!r}." + ) + if self.token_in_async_flags and self.token_in_flag_batch != 1: + raise ValueError("Asynchronous token-in flags require token_in_flag_batch=1.") + if not 1 <= self.token_in_flag_batch <= 32: + raise ValueError(f"token_in_flag_batch must be in [1, 32], got {self.token_in_flag_batch}.") + if ( + isinstance(self.transfer_warp_count, bool) + or not isinstance(self.transfer_warp_count, int) + or not 1 <= self.transfer_warp_count <= 32 + ): + raise ValueError( + "token_in_transfer_warp_count must be in [1, 32], got " + f"{self.transfer_warp_count}." + ) + if self.tokens_per_fc1_ready_slot % self.token_padding_block != 0: + raise ValueError("tokens_per_fc1_ready_slot must be divisible by token_padding_block.") + _pad_lo = min(self.sf_padding_block, self.token_padding_block) + _pad_hi = max(self.sf_padding_block, self.token_padding_block) + if _pad_hi % _pad_lo != 0: + raise ValueError("sf_padding_block and token_padding_block must be power-of-two multiples of each other.") + if self.token_back_enabled and self.fc2_done_signals_per_token_tile <= 0: + raise ValueError("fc2_done_signals_per_token_tile must be positive when token-back is enabled.") + if self.reduce_topk_in_kernel and self.combine_format.act_dtype is not cutlass.BFloat16: + raise ValueError("In-kernel top-k reduction requires BF16 combine data.") + element_block = self.activation_sf_vector_size * 4 + if self.hidden_size % element_block != 0: + raise ValueError(f"{self.quant_kind} requires hidden_size divisible by {element_block}.") + if self.sf_padding_block % 128 != 0: + raise ValueError("sf_padding_block must be a multiple of 128.") + + @property + def experts_per_rank(self) -> int: + return self.expert_count // self.world_size + + @property + def activation_dtype(self) -> type: + return _quant_spec[self.quant_kind][0] + + @property + def activation_sf_dtype(self) -> type: + return _quant_spec[self.quant_kind][1] + + @property + def activation_sf_vector_size(self) -> int: + return _quant_spec[self.quant_kind][2] + + @property + def bytes_per_token(self) -> int: + return self.hidden_size * int(self.activation_dtype.width) // 8 + + @property + def activation_sf_hidden_padded(self) -> int: + valid_hidden = self.hidden_size // self.activation_sf_vector_size + elements_per_16_bytes = 128 // int(self.activation_sf_dtype.width) + return int(round_up(valid_hidden, elements_per_16_bytes)) + + @property + def combine_sf_hidden_padded(self) -> int: + if not self.combine_format.is_quantized: + return 0 + valid_hidden = self.hidden_size // self.combine_format.scale_block + elements_per_16_bytes = 128 // int(self.combine_format.scale_dtype.width) + return int(round_up(valid_hidden, elements_per_16_bytes)) + + @property + def token_back_push_data(self) -> bool: + return self.token_back_mode != "epi_warps" # nosec B105 + + @property + def token_back_push_sf(self) -> bool: + return self.combine_format.is_quantized + + @property + def token_back_enabled(self) -> bool: + return self.token_back_push_data or self.token_back_push_sf + + @property + def worst_case_token_count(self) -> int: + return self._router.worst_case_token_count + + @property + def worst_case_sf_token_count(self) -> int: + # The SF pool has to be at least as tall as the data pool: once + # token_padding_block may exceed sf_padding_block, an expert's data + # cursor can walk past the SF-capacity-derived row count. + sf_capacity = self._router.receive_capacity(self.sf_padding_block).padded_route_count + return max(sf_capacity, self.worst_case_token_count) + + @property + def max_fc1_ready_slot_count(self) -> int: + return ( + self._router.receive_capacity(self.tokens_per_fc1_ready_slot).padded_route_count + // self.tokens_per_fc1_ready_slot + ) + + @property + def router_smem_workspace(self) -> SmemWorkspace: + return self._router.router_smem_workspace + + @property + def expert_count_padded(self) -> int: + return self._router.expert_count_padded + + @property + def expert_count_with_trash(self) -> int: + return self._router.expert_count_with_trash + + @property + def router_elements_per_lane(self) -> int: + return self._router.router_elements_per_lane + + @property + def router_warps_per_cta(self) -> int: + return self._router.router_warps_per_cta + + @property + def router_data_cta_count(self) -> int: + return self._router.router_data_cta_count + + @property + def router_tokens_per_cta(self) -> int: + return self._router.router_tokens_per_cta + + @property + def router_push_cta_count(self) -> int: + return self._router.router_push_cta_count + + @property + def router_grid_cta_count(self) -> int: + return self._router.router_grid_cta_count + + @property + def metadata_ready_target(self) -> int: + return self._router.metadata_ready_target + + def register_device_workspace(self, workspace: DeviceWorkspace) -> None: + self._router.register_device_workspace(workspace) + self._register_main_workspace(workspace) + self._nvlink_barrier.register_device_workspace(workspace) + + @cute.jit + def launch_router( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + overflow_flag: cute.Tensor, + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper_host, + device_workspace: DeviceWorkspace, + stream: cuda.CUstream, + ) -> None: + self._router.launch_router( + topk_indices, + topk_scores, + overflow_flag, + local_rank, + local_workspace, + shared_workspace, + peer_rank_ptr_mapper_host, + device_workspace, + stream, + ) + + @cute.jit + def local_expert_sizes(self, device_workspace: DeviceWorkspace, local_rank: Int32) -> cute.Tensor: + return self._router.local_expert_sizes(device_workspace, local_rank) + + @cute.jit + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + self._router.wait_for_sizes_ready(device_workspace, sleep_cycles) + + @cute.jit + def token_src_metadata_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return self._router.token_src_metadata_tensor(device_workspace) + + @cute.jit + def fc1_topk_scores_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + return self._router.fc1_topk_scores_tensor(device_workspace) + + @cute.jit + def assign_device_members( + self, + *, + device_workspace: DeviceWorkspace, + token_comm_args: TokenCommArgs, + local_rank: Int32, + linear_cta_idx: Int32, + ) -> None: + self._device_workspace = device_workspace + self._token_comm_args = token_comm_args + self._local_rank = local_rank + self._linear_cta_idx = linear_cta_idx + thread_idx, _, _ = cute.arch.thread_idx() + transfer_thread_idx = thread_idx % Int32(self.transfer_thread_count) + self._transfer_warp_idx = cute.arch.make_warp_uniform(transfer_thread_idx // Int32(32)) + self._lane_idx = transfer_thread_idx % Int32(32) + self._nvlink_barrier.assign_device_members(device_workspace, token_comm_args.peer_rank_ptr_mapper) + + def remove_device_members(self) -> None: + self._nvlink_barrier.remove_device_members() + self._device_workspace = None + self._token_comm_args = None + self._local_rank = None + self._linear_cta_idx = None + self._transfer_warp_idx = None + self._lane_idx = None + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "TokenCommDeterministic": + if values: + raise ValueError("TokenCommDeterministic carries no MLIR values.") + return self + + def _register_main_workspace(self, workspace: DeviceWorkspace) -> None: + workspace.register( + self.fc1_ready_region, + cutlass.Int32, + (self.max_fc1_ready_slot_count,), + buffer_space="local", + reset="tail_reset", + ) + activation_element_count = self.worst_case_token_count * self.hidden_size + workspace.register( + self.fc1_activation_region, + self.activation_dtype, + (activation_element_count,), + buffer_space="local", + byte_alignment=128, + ) + activation_sf_element_count = self.worst_case_sf_token_count * self.activation_sf_hidden_padded + workspace.register( + self.fc1_activation_sf_region, + self.activation_sf_dtype, + (activation_sf_element_count,), + buffer_space="local", + byte_alignment=128, + ) + if self.token_back_enabled: + workspace.register( + self.fc2_done_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local", reset="tail_reset" + ) + if self.token_back_push_data: + fc2_element_count = self.worst_case_token_count * self.hidden_size + workspace.register( + self.fc2_activation_region, + self.combine_format.act_dtype, + (fc2_element_count,), + buffer_space="local", + byte_alignment=128, + ) + if self.token_back_push_sf: + fc2_sf_element_count = self.worst_case_token_count * self.combine_sf_hidden_padded + workspace.register( + self.fc2_activation_sf_region, + self.combine_format.scale_dtype, + (fc2_sf_element_count,), + buffer_space="local", + byte_alignment=128, + ) + if self.token_back_enabled and self.token_back_schedule_mode == "atomic_counter": # nosec B105 + workspace.register( + self.token_back_schedule_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + if not self.reduce_topk_in_kernel: + workspace.register( + self.pre_reduced_activation_region, + self.combine_format.act_dtype, + (self.max_tokens_per_rank, self.topk, self.hidden_size), + buffer_space="shared", + mem_order=(2, 1, 0), + byte_alignment=128, + ) + if self.combine_format.is_quantized: + workspace.register( + self.pre_reduced_activation_sf_region, + self.combine_format.scale_dtype, + (self.max_tokens_per_rank, self.topk, self.combine_sf_hidden_padded), + buffer_space="shared", + mem_order=(2, 1, 0), + byte_alignment=128, + ) + + def register_smem_regions(self, workspace: SmemWorkspace) -> None: + workspace.register_mbarrier(self.token_in_mbarrier_region, self.transfer_warp_count) + if self.token_back_enabled: + workspace.register_mbarrier(self.token_back_mbarrier_region, self.transfer_warp_count) + workspace.register_tensor( + self.expert_sizes_smem_region, cutlass.Int32, (self.experts_per_rank,), byte_alignment=16 + ) + transfer_overlay = workspace.create_overlay("nvlink.token_comm.main_smem.transfer") + token_in_lifetime = transfer_overlay.add_lifetime("token_in") + token_in_lifetime.register_tensor( + self.token_in_activation_smem_region, + self.activation_dtype, + (self.transfer_warp_count, self.hidden_size), + byte_alignment=16, + ) + token_in_lifetime.register_tensor( + self.token_in_sf_smem_region, + self.activation_sf_dtype, + (self.transfer_warp_count, (self.activation_sf_vector_size, self.activation_sf_hidden_padded)), + stride=(self.activation_sf_hidden_padded, (0, 1)), + byte_alignment=16, + ) + if not self.token_back_enabled: + return + + if self.token_back_mode == "standalone_warps": # nosec B105 + token_back_lifetime = workspace.create_overlay( + "nvlink.token_comm.main_smem.standalone_token_back" + ).add_lifetime("token_back") + else: + token_back_lifetime = transfer_overlay.add_lifetime("token_back") + + if self.token_back_mode == "standalone_warps": # nosec B105 + available_bytes_per_warp = self.standalone_chunk_bytes + else: + activation_bytes = self.bytes_per_token + sf_bytes = self.activation_sf_hidden_padded * int(self.activation_sf_dtype.width) // 8 + available_bytes_per_warp = activation_bytes + sf_bytes + + if self.token_back_push_data: + bytes_per_output_token = self.hidden_size * int(self.combine_format.act_dtype.width) // 8 + if self.token_back_mode == "standalone_warps": # nosec B105 + chunk_bytes = self.standalone_chunk_bytes + elif available_bytes_per_warp < bytes_per_output_token: + chunk_bytes = self.bytes_per_token + else: + chunk_bytes = bytes_per_output_token + if self.token_back_mode != "standalone_warps" and bytes_per_output_token % chunk_bytes != 0: # nosec B105 + raise ValueError("Token-back data chunk bytes must divide one row.") + chunk_elements = chunk_bytes * 8 // int(self.combine_format.act_dtype.width) + token_back_lifetime.register_tensor( + self.token_back_activation_smem_region, + self.combine_format.act_dtype, + (self.transfer_warp_count, chunk_elements), + byte_alignment=16, + ) + if self.token_back_push_sf: + sf_row_bytes = self.combine_sf_hidden_padded * int(self.combine_format.scale_dtype.width) // 8 + if sf_row_bytes > available_bytes_per_warp: + raise ValueError("Token-back scale row exceeds its per-warp stage.") + token_back_lifetime.register_tensor( + self.token_back_sf_smem_region, + self.combine_format.scale_dtype, + (self.transfer_warp_count, (self.combine_format.scale_block, self.combine_sf_hidden_padded)), + stride=(self.combine_sf_hidden_padded, (0, 1)), + byte_alignment=16, + ) + + @cute.jit + def fc1_ready_counter_pointer(self, device_workspace: DeviceWorkspace) -> cute.Pointer: + return device_workspace.ptr(self.fc1_ready_region) + + @cute.jit + def fc1_activation_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return cute.make_tensor( + device_workspace.ptr(self.fc1_activation_region), + cute.make_layout((self.worst_case_token_count, self.hidden_size), stride=(self.hidden_size, 1)), + ) + + @cute.jit + def fc1_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + layout = tile_atom_to_shape_SF( + (self.worst_case_sf_token_count, self.hidden_size, 1), self.activation_sf_vector_size + ) + return cute.make_tensor(device_workspace.ptr(self.fc1_activation_sf_region), cute.select(layout, mode=[0, 1])) + + @cute.jit + def fc2_done_counter_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.token_back_enabled): + return None + return device_workspace.tensor(self.fc2_done_region) + + @cute.jit + def fc2_activation_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.token_back_push_data): + return None + return cute.make_tensor( + device_workspace.ptr(self.fc2_activation_region), + cute.make_layout( + (self.worst_case_token_count, 1, self.hidden_size), stride=(self.hidden_size, self.hidden_size, 1) + ), + ) + + @cute.jit + def fc2_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.token_back_push_sf): + return None + return cute.make_tensor( + device_workspace.ptr(self.fc2_activation_sf_region), + cute.make_layout( + ( + self.worst_case_token_count, + 1, + (self.combine_format.scale_block, self.hidden_size // self.combine_format.scale_block), + ), + stride=(self.combine_sf_hidden_padded, self.combine_sf_hidden_padded, (0, 1)), + ), + ) + + @cute.jit + def pre_reduced_activation_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + """Return the source-domain combine plane when top-k reduction is separate.""" + if cutlass.const_expr(self.reduce_topk_in_kernel): + return None + return device_workspace.tensor(self.pre_reduced_activation_region) + + @cute.jit + def pre_reduced_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + """Return the source-domain combine scale plane when one is required.""" + if cutlass.const_expr(self.reduce_topk_in_kernel or not self.combine_format.is_quantized): + return None + return device_workspace.tensor(self.pre_reduced_activation_sf_region) + + @cute.jit + def _load_input_dedup_route_batch( + self, routes: cute.Pointer, source_global_token: Int32, route_batch: Int32, + ) -> Int64: + """Keep one route per lane; predicate the final partial batch.""" + route_slot = route_batch * Int32(32) + self._lane_idx + record = Int64(0) + if route_slot < Int32(self.topk): + record = cute.arch.load( + routes + source_global_token * Int32(self.topk) + route_slot, Int64 + ) + return record + + @cute.jit + def _token_in_deduplicated(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: + """One remote TMA payload per source token/receiver group, original row outputs. + + Router writes a direct slot index before its original metadata publication. + Each slot has one source-router writer. Full tail drain precedes index reset. + The original resident communication warps guarantee progress; no new kernel, + cross-CTA waiting, mathematical operation, or reduced output is introduced. + """ + iket.range_push("token_in.deduplicated_payload") + lane_idx = self._lane_idx + transfer_warp_idx = self._transfer_warp_idx + route_batch_count = ceil_div(self.topk, 32) + global_warp_idx = self._linear_cta_idx * Int32(self.transfer_warp_count) + transfer_warp_idx + global_warp_count = Int32(self.promised_launchable_sm_count * self.transfer_warp_count) + routes = self._device_workspace.ptr(self._router.input_dedup_routes_region) + stage = smem_workspace.tensor(self.token_in_activation_smem_region, smem_base)[transfer_warp_idx, None] + sf_stage = smem_workspace.tensor(self.token_in_sf_smem_region, smem_base)[transfer_warp_idx, (None, None)] + barrier = smem_workspace.ptr(self.token_in_mbarrier_region, smem_base) + transfer_warp_idx + if lane_idx == Int32(0): + cute.arch.mbarrier_init(barrier, 1) + cute.arch.sync_warp() + activation_bytes = cute.cosize(stage) * int(self.activation_dtype.width) // 8 + sf_bytes = cute.cosize(sf_stage) * int(self.activation_sf_dtype.width) // 8 + source_sf_values = cute.slice_(sf_stage, (0, None)) + source_sf_vectors = cute.zipped_divide(source_sf_values, (4,)) + sf_atom = _copy_atom(self.activation_sf_dtype, 4 * int(self.activation_sf_dtype.width)) + destination_sf = self.fc1_activation_sf_tensor(self._device_workspace) + destination_data = self._device_workspace.ptr(self.fc1_activation_region) + ready = self._device_workspace.ptr(self.fc1_ready_region) + flags = make_flag_batch_tracker(use_async=self.token_in_async_flags, flag_address=Int64(0), + accumulated_flags=Int32(0), phase=Int32(0), thread_idx=lane_idx) + pull_phase = Int32(0) + # Traverse the same dense, expert-major valid rows as original + # token_in; physical pool padding never becomes a work item. + owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) + pool_bases = self._router.pool_expert_base_tensor(self._device_workspace) + metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + work_limit = Int32(0) + for expert_sum in cutlass.range_constexpr(self.experts_per_rank): + work_limit = work_limit + owned_sizes[Int32(expert_sum)] + current_expert = Int32(0) + expert_valid_begin = Int32(0) + expert_valid_end = owned_sizes[Int32(0)] + work = global_warp_idx + while work < work_limit: + # work < sum(sizes) proves a nonempty expert remains, including + # leading/interior zero-size experts and the final partial wave. + while work >= expert_valid_end: + expert_valid_begin = expert_valid_end + current_expert = current_expert + Int32(1) + expert_valid_end = expert_valid_end + owned_sizes[current_expert] + pool_row = pool_bases[current_expert] + work - expert_valid_begin + source_rank_from_row = Int32(0) + source_token_from_row = Int32(0) + if lane_idx == Int32(0): + metadata = TokenSrcMetadata.load( + metadata_pointer.toint() + Int64(pool_row) * Int64(TokenSrcMetadata.nbytes) + ) + source_rank_from_row = metadata.src_rank + source_token_from_row = metadata.src_token + source_rank_from_row = Int32(cute.arch.shuffle_sync(source_rank_from_row, Int32(0))) + source_token_from_row = Int32(cute.arch.shuffle_sync(source_token_from_row, Int32(0))) + source_global_token = source_rank_from_row * Int32(self.max_tokens_per_rank) + source_token_from_row + # Every route owns a distinct destination row. Select the smallest + # row across ALL batches, preserving expert-major ownership without + # packing a local expert ID. Repeated experts still have distinct rows. + record = self._load_input_dedup_route_batch(routes, source_global_token, Int32(0)) + route_row = Int32(record & Int64(0xFFFFF)) + earlier_owner = (route_row != Int32(0)) & (route_row < pool_row + Int32(1)) + has_earlier_owner = cute.arch.vote_ballot_sync(earlier_owner) + if cutlass.const_expr(self.topk > 32): + for route_batch in cutlass.range(1, route_batch_count, unroll=1): + other = self._load_input_dedup_route_batch(routes, source_global_token, route_batch) + other_row = Int32(other & Int64(0xFFFFF)) + earlier_owner = (other_row != Int32(0)) & (other_row < pool_row + Int32(1)) + has_earlier_owner = has_earlier_owner | cute.arch.vote_ballot_sync(earlier_owner) + # The current valid row guarantees at least one local route, even + # when this receiver's first matching route is in a later batch. + if has_earlier_owner == Int32(0): + source_rank = source_global_token // Int32(self.max_tokens_per_rank) + source_token = source_global_token % Int32(self.max_tokens_per_rank) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), source_rank, Int64(0)) + source_data_address = self._token_comm_args.activation.iterator.toint() + peer_offset + Int64(source_token) * Int64(self.bytes_per_token) + remote_sf = cute.make_tensor(cute.make_ptr( + self._token_comm_args.activation_sf.dtype, + self._token_comm_args.activation_sf.iterator.toint() + peer_offset, + AddressSpace.gmem, assumed_align=self._token_comm_args.activation_sf.iterator.max_alignment), + self._token_comm_args.activation_sf.layout) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(barrier, Int32(activation_bytes + sf_bytes)) + tma_load_1d(stage.iterator, source_data_address, barrier, Int32(activation_bytes)) + tma_load_1d(sf_stage.iterator, remote_sf[Int64(source_token), None].iterator, barrier, Int32(sf_bytes)) + cute.arch.sync_warp() + cute.arch.mbarrier_wait(barrier, pull_phase) + for route_batch in cutlass.range(route_batch_count, unroll=1): + batch_record = record + if cutlass.const_expr(self.topk > 32): + if route_batch != Int32(0): + batch_record = self._load_input_dedup_route_batch( + routes, source_global_token, route_batch + ) + for slot in cutlass.range_constexpr(min(self.topk, 32)): + lo = cute.arch.shuffle_sync(Int32(batch_record & Int64(0xFFFFFFFF)), Int32(slot)) + hi = cute.arch.shuffle_sync(Int32(batch_record >> Int64(32)), Int32(slot)) + route = (Int64(hi) << Int64(32)) | (Int64(lo) & Int64(0xFFFFFFFF)) + if route != Int64(0): + destination_row = Int32(route & Int64(0xFFFFF)) - Int32(1) + sf_row = Int32((route >> Int64(20)) & Int64(0xFFFFF)) + dst = cute.make_ptr(self.activation_dtype, + destination_data.toint() + Int64(destination_row) * Int64(self.bytes_per_token), + AddressSpace.gmem, assumed_align=16) + with cute.arch.elect_one(): + cp_async_bulk_s2g(dst, stage.iterator, Int32(activation_bytes)) + cute.arch.sync_warp() + sf_row_view = destination_sf[Int64(sf_row), ((None, None), None)] + sf_values = cute.group_modes(cute.slice_(sf_row_view, (0, None, None)), 0, 2) + sf_vectors = cute.zipped_divide(sf_values, (4,)) + for sf_round in cutlass.range_constexpr(ceil_div(cute.size(sf_vectors, mode=[1]), 32)): + vector = Int32(sf_round * 32) + lane_idx + if vector < Int32(cute.size(sf_vectors, mode=[1])): + cute.copy(sf_atom, source_sf_vectors[None, vector], sf_vectors[None, vector]) + # Every batch has issued its stores before any ready is + # published. All destinations are distinct pool rows. Their S2G + # copies may share one bulk async-group without write-order + # dependencies. Keep both source stages immutable until the + # full wait below; .read alone is insufficient for ready. + cute.arch.sync_warp() + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0) + # Join every bulk issuer and every generic SF writer before + # the elected lane publishes any destination's release flag. + cute.arch.sync_warp() + for route_batch in cutlass.range(route_batch_count, unroll=1): + batch_record = record + if cutlass.const_expr(self.topk > 32): + if route_batch != Int32(0): + batch_record = self._load_input_dedup_route_batch( + routes, source_global_token, route_batch + ) + for slot in cutlass.range_constexpr(min(self.topk, 32)): + lo = cute.arch.shuffle_sync(Int32(batch_record & Int64(0xFFFFFFFF)), Int32(slot)) + hi = cute.arch.shuffle_sync(Int32(batch_record >> Int64(32)), Int32(slot)) + route = (Int64(hi) << Int64(32)) | (Int64(lo) & Int64(0xFFFFFFFF)) + if route != Int64(0): + ready_slot = Int32((route >> Int64(40)) & Int64(0xFFFFF)) + flags = flags.accumulate( + Int32(0), self.token_in_flag_batch, (ready + ready_slot).toint() + ) + cute.arch.sync_warp() + pull_phase = pull_phase ^ Int32(1) + work = work + global_warp_count + flags.fire() + cute.arch.sync_warp() + iket.range_pop() + + @cute.jit + def token_in( + self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer, + ) -> None: + """Wait for pushed metadata, then pull activation payloads into local pools.""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + global_warp_idx = self._linear_cta_idx * Int32(self.transfer_warp_count) + transfer_warp_idx + global_warp_count = Int32(self.promised_launchable_sm_count * self.transfer_warp_count) + + sizes = self._router.sizes_tensor(self._device_workspace) + pool_expert_bases = self._router.pool_expert_base_tensor(self._device_workspace) + token_metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + + iket.range_push("token_in.wait_sizes_ready") + self._router.wait_for_sizes_ready(self._device_workspace) + iket.range_pop() + iket.range_push("token_in.stage_sizes") + owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) + owner_expert_begin = self._local_rank * Int32(self.experts_per_rank) + source_sizes = cute.make_tensor(sizes.iterator + owner_expert_begin, cute.make_layout((self.experts_per_rank,))) + copy_elements = 4 if self.experts_per_rank % 4 == 0 else 1 + source_size_vectors = cute.zipped_divide( + (_mark_alignment(source_sizes, 16) if cutlass.const_expr(copy_elements == 4) else source_sizes), + (copy_elements,), + ) + destination_size_vectors = cute.zipped_divide(owned_sizes, (copy_elements,)) + size_vector_count = cute.size(destination_size_vectors, mode=[1]) + size_copy_atom = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp( + cache_mode=cute.nvgpu.LoadCacheMode.GLOBAL if copy_elements == 4 else cute.nvgpu.LoadCacheMode.ALWAYS + ), + cutlass.Int32, + num_bits_per_copy=copy_elements * 32, + ) + size_copy_rounds = ceil_div(size_vector_count, self.transfer_thread_count) + transfer_thread_idx = transfer_warp_idx * Int32(32) + lane_idx + for size_copy_round in cutlass.range_constexpr(size_copy_rounds): + vector_idx = Int32(size_copy_round * self.transfer_thread_count) + transfer_thread_idx + if vector_idx < Int32(size_vector_count): + cute.copy( + size_copy_atom, source_size_vectors[None, vector_idx], destination_size_vectors[None, vector_idx] + ) + cute.arch.cp_async_commit_group() + iket.range_pop() + + iket.range_push("token_in.wait_metadata_ready") + self._router.wait_for_metadata_ready(self._device_workspace) + iket.range_pop() + + cute.arch.cp_async_wait_group(0) + iket.range_push("token_in.size_barrier") + token_in_size_barrier = pipeline.NamedBarrier( + barrier_id=self.token_in_size_barrier_id, num_threads=self.transfer_thread_count + ) + token_in_size_barrier.arrive_and_wait() + iket.range_pop() + if cutlass.const_expr(self.token_back_mode == "standalone_warps"): # nosec B105 + sizes_ready_barrier = pipeline.NamedBarrier( + barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count + ) + sizes_ready_barrier.arrive() + + if cutlass.const_expr(self.input_dedup_mode != "off"): + self._token_in_deduplicated(smem_workspace, smem_base) + if cutlass.const_expr(self.input_dedup_mode == "off"): + iket.range_push("token_in.pull_payload") + token_in_mbarriers = smem_workspace.ptr(self.token_in_mbarrier_region, smem_base) + token_in_activation = smem_workspace.tensor(self.token_in_activation_smem_region, smem_base) + token_in_sf = smem_workspace.tensor(self.token_in_sf_smem_region, smem_base) + warp_mbarrier = token_in_mbarriers + transfer_warp_idx + warp_activation_stage = token_in_activation[transfer_warp_idx, None] + warp_sf_stage = token_in_sf[transfer_warp_idx, (None, None)] + if lane_idx == Int32(0): + cute.arch.mbarrier_init(warp_mbarrier, 1) + cute.arch.sync_warp() + + fc1_activation_pointer = self._device_workspace.ptr(self.fc1_activation_region) + fc1_activation_sf = self.fc1_activation_sf_tensor(self._device_workspace) + fc1_ready_counter = self._device_workspace.ptr(self.fc1_ready_region) + activation_bytes = cute.cosize(warp_activation_stage) * int(self.activation_dtype.width) // 8 + activation_sf_bytes = cute.cosize(warp_sf_stage) * int(self.activation_sf_dtype.width) // 8 + sf_copy_elements = 4 + source_sf_values = cute.slice_(warp_sf_stage, (0, None)) + source_sf_vectors = cute.zipped_divide(source_sf_values, (sf_copy_elements,)) + sf_copy_atom = _copy_atom(self.activation_sf_dtype, sf_copy_elements * int(self.activation_sf_dtype.width)) + + next_dense_token = global_warp_idx + expert_valid_begin = Int32(0) + expert_sf_begin = Int32(0) + expert_ready_slot_begin = Int32(0) + pull_phase = Int32(0) + flag_tracker = make_flag_batch_tracker( + use_async=self.token_in_async_flags, + flag_address=Int64(0), + accumulated_flags=Int32(0), + phase=Int32(0), + thread_idx=lane_idx, + ) + + local_expert = Int32(0) + while local_expert < Int32(self.experts_per_rank): + expert_token_count = owned_sizes[local_expert] + pool_expert_base = pool_expert_bases[local_expert] + expert_valid_end = expert_valid_begin + expert_token_count + pull_count = Int32(0) + if next_dense_token < expert_valid_end: + pull_count = (expert_valid_end - next_dense_token + global_warp_count - Int32(1)) // global_warp_count + + for pull_round in cutlass.range(pull_count, unroll=1): + dense_token_idx = next_dense_token + Int32(pull_round) * global_warp_count + token_in_expert = dense_token_idx - expert_valid_begin + pool_token_idx = pool_expert_base + token_in_expert + sf_token_idx = expert_sf_begin + token_in_expert + ready_slot_idx = expert_ready_slot_begin + token_in_expert // Int32(self.tokens_per_fc1_ready_slot) + + metadata = TokenSrcMetadata.load( + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) + ) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), metadata.src_rank, Int64(0)) + remote_activation_address = ( + self._token_comm_args.activation.iterator.toint() + + peer_offset + + Int64(metadata.src_token) * Int64(self.bytes_per_token) + ) + remote_sf = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.activation_sf.dtype, + self._token_comm_args.activation_sf.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.activation_sf.iterator.max_alignment), + ), + self._token_comm_args.activation_sf.layout, + ) + remote_sf_row = remote_sf[Int64(metadata.src_token), None] + + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + warp_mbarrier, Int32(activation_bytes + activation_sf_bytes) + ) + tma_load_1d( + warp_activation_stage.iterator, + Int64(remote_activation_address), + warp_mbarrier, + Int32(activation_bytes), + ) + tma_load_1d( + warp_sf_stage.iterator, remote_sf_row.iterator, warp_mbarrier, Int32(activation_sf_bytes) + ) + cute.arch.sync_warp() + cute.arch.mbarrier_wait(warp_mbarrier, pull_phase) + + destination_activation = cute.make_ptr( + self.activation_dtype, + fc1_activation_pointer.toint() + Int64(pool_token_idx) * Int64(self.bytes_per_token), + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_activation, warp_activation_stage.iterator, Int32(activation_bytes)) + cute.arch.sync_warp() + cute.arch.cp_async_bulk_commit_group() + + destination_sf_row = fc1_activation_sf[Int64(sf_token_idx), ((None, None), None)] + destination_sf_values = cute.slice_(destination_sf_row, (0, None, None)) + destination_sf_values = cute.group_modes(destination_sf_values, 0, 2) + destination_sf_vectors = cute.zipped_divide(destination_sf_values, (sf_copy_elements,)) + sf_vector_count = cute.size(destination_sf_vectors, mode=[1]) + for sf_round in cutlass.range_constexpr(ceil_div(sf_vector_count, 32)): + sf_vector_idx = Int32(sf_round * 32) + lane_idx + if sf_vector_idx < Int32(sf_vector_count): + cute.copy( + sf_copy_atom, + source_sf_vectors[None, sf_vector_idx], + destination_sf_vectors[None, sf_vector_idx], + ) + + if cutlass.const_expr(self.token_in_flag_batch == 1): + cute.arch.cp_async_bulk_wait_group(0) + else: + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.sync_warp() + ready_address = (fc1_ready_counter + ready_slot_idx).toint() + flag_tracker = flag_tracker.accumulate(Int32(0), self.token_in_flag_batch, ready_address) + cute.arch.sync_warp() + pull_phase = pull_phase ^ Int32(1) + + next_dense_token = next_dense_token + pull_count * global_warp_count + expert_valid_begin = expert_valid_end + expert_sf_begin = expert_sf_begin + ( + (expert_token_count + Int32(self.sf_padding_block - 1)) // Int32(self.sf_padding_block) + ) * Int32(self.sf_padding_block) + expert_ready_slot_begin = expert_ready_slot_begin + ( + (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) + // Int32(self.tokens_per_fc1_ready_slot) + ) + local_expert = local_expert + Int32(1) + + flag_tracker.fire() + cute.arch.sync_warp() + iket.range_pop() + if cutlass.const_expr(self.token_back_enabled and self.token_back_mode != "standalone_warps"): # nosec B105 + iket.range_push("token_in.transfer_barrier") + transfer_lifetime_barrier = pipeline.NamedBarrier( + barrier_id=self.transfer_lifetime_barrier_id, num_threads=self.transfer_thread_count + ) + transfer_lifetime_barrier.arrive_and_wait() + iket.range_pop() + + @cute.jit + def _stateless_pace(self, reference_window: Int32, current_window: Int32) -> None: + sleep_cycles = Int32(0) + if current_window < reference_window: + sleep_cycles = reference_window - current_window + elif current_window > reference_window: + sleep_cycles = cutlass.min(current_window - reference_window, Int32(self.standalone_max_backoff_cycles)) + if sleep_cycles > Int32(0): + nanosleep(sleep_cycles) + + @cute.jit + def _adaptive_pace(self, average_window: Int32, current_window: Int32, low_window: int, high_window: int) -> Int32: + sleep_cycles = Int32(0) + if current_window > average_window: + average_window = average_window + ((current_window - average_window + Int32(3)) // Int32(4)) + sleep_cycles = current_window - average_window + if sleep_cycles > Int32(high_window): + sleep_cycles = Int32(high_window) + else: + average_window = average_window - ((average_window - current_window + Int32(3)) // Int32(4)) + sleep_cycles = average_window - current_window + if sleep_cycles > Int32(self.adaptive_minimum_sleep_cycles): + nanosleep(sleep_cycles) + if average_window > Int32(high_window): + average_window = Int32(high_window) + if average_window < Int32(low_window): + average_window = Int32(low_window) + return average_window + + @cute.jit + def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: + """Push completed FC2 data and scale rows to source ranks.""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + if cutlass.const_expr(not self.token_back_enabled): + return + if cutlass.const_expr( + self.combine_format.is_quantized and self._token_comm_args.pre_reduced_activation_sf is None + ): + raise ValueError("Quantized token-back requires a scale destination.") + + global_worker_idx = self._linear_cta_idx * Int32(self.transfer_warp_count) + transfer_warp_idx + global_worker_count = Int32(self.promised_launchable_sm_count * self.transfer_warp_count) + token_back_mbarriers = smem_workspace.ptr(self.token_back_mbarrier_region, smem_base) + worker_mbarrier = token_back_mbarriers + transfer_warp_idx + if cutlass.const_expr(self.token_back_push_data): + token_back_activation = smem_workspace.tensor(self.token_back_activation_smem_region, smem_base) + worker_activation_stage = token_back_activation[transfer_warp_idx, None] + activation_chunk_bytes = ( + cute.cosize(worker_activation_stage) * int(self.combine_format.act_dtype.width) // 8 + ) + if cutlass.const_expr(self.token_back_push_sf): + token_back_sf = smem_workspace.tensor(self.token_back_sf_smem_region, smem_base) + worker_sf_stage = token_back_sf[transfer_warp_idx, (None, None)] + sf_chunk_bytes = cute.cosize(worker_sf_stage) * int(self.combine_format.scale_dtype.width) // 8 + if lane_idx == Int32(0): + cute.arch.mbarrier_init(worker_mbarrier, 1) + cute.arch.sync_warp() + + owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) + if cutlass.const_expr(self.token_back_mode == "standalone_warps"): # nosec B105 + sizes_ready_barrier = pipeline.NamedBarrier( + barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count + ) + sizes_ready_barrier.arrive_and_wait() + + pool_expert_bases = self._router.pool_expert_base_tensor(self._device_workspace) + token_metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + fc2_done = self._device_workspace.ptr(self.fc2_done_region) + if cutlass.const_expr(self.token_back_push_data): + fc2_activation_pointer = self._device_workspace.ptr(self.fc2_activation_region) + output_token_bytes = self.hidden_size * int(self.combine_format.act_dtype.width) // 8 + activation_chunk_count = ceil_div(output_token_bytes, activation_chunk_bytes) + data_window_unit = ceil_div(activation_chunk_bytes * 2, 3) + reuse_data_pacing_enabled = ( + self.token_back_mode == "reuse_dispatch_warps" and data_window_unit > self.minimum_pacing_window_cycles # nosec B105 + ) + # Preserve the empirical low:initial:high ratio of 1:2.5:5. + data_average_window = Int32(data_window_unit) + data_low_window = data_window_unit * 2 // 5 + data_high_window = data_window_unit * 2 + print(f"[{data_window_unit}, {data_low_window}, {data_high_window}]") + if cutlass.const_expr(self.token_back_push_sf): + fc2_sf_pointer = self._device_workspace.ptr(self.fc2_activation_sf_region) + output_sf_bytes = self.combine_sf_hidden_padded * int(self.combine_format.scale_dtype.width) // 8 + sf_chunk_count = ceil_div(output_sf_bytes, sf_chunk_bytes) + sf_window_unit = ceil_div(sf_chunk_bytes * 2, 3) + reuse_sf_pacing_enabled = ( + self.token_back_mode == "reuse_dispatch_warps" and sf_window_unit > self.minimum_pacing_window_cycles # nosec B105 + ) + sf_average_window = Int32(sf_window_unit) + sf_low_window = sf_window_unit * 2 // 5 + sf_high_window = sf_window_unit * 2 + + next_dense_token = global_worker_idx - global_worker_count + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 + next_dense_token = Int32(0) + next_dense_token = self.next_token(next_dense_token) + expert_valid_begin = Int32(0) + transfer_phase = Int32(0) + + iket.range_push("token_back.work") + local_expert = Int32(0) + while local_expert < Int32(self.experts_per_rank): + expert_token_count = owned_sizes[local_expert] + expert_valid_end = expert_valid_begin + expert_token_count + if next_dense_token < expert_valid_end: + token_tile_count = (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) // Int32( + self.tokens_per_fc1_ready_slot + ) + completion_target = token_tile_count * Int32(self.fc2_done_signals_per_token_tile) + iket.range_push("token_back.wait_fc2") + while cute.arch.load(fc2_done + local_expert, Int32, sem="acquire", scope="gpu") < completion_target: + nanosleep(500) + iket.range_pop() + + while next_dense_token < expert_valid_end: + token_in_expert = next_dense_token - expert_valid_begin + pool_token_idx = pool_expert_bases[local_expert] + token_in_expert + metadata = TokenSrcMetadata.load( + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) + ) + destination_topk = metadata.src_topk + if cutlass.const_expr(self.reduce_topk_in_kernel): + destination_topk = Int32(0) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), metadata.src_rank, Int64(0)) + is_remote_token = metadata.src_rank != self._local_rank + + if cutlass.const_expr(self.token_back_push_data): + iket.range_push("token_back.push_data") + local_activation_address = fc2_activation_pointer.toint() + Int64(pool_token_idx) * Int64( + output_token_bytes + ) + remote_activation = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.pre_reduced_activation.dtype, + self._token_comm_args.pre_reduced_activation.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.pre_reduced_activation.iterator.max_alignment), + ), + self._token_comm_args.pre_reduced_activation.layout, + ) + destination_row = remote_activation[Int64(metadata.src_token), destination_topk, None] + for chunk_idx in cutlass.range_constexpr(activation_chunk_count): + chunk_byte_offset = Int64(chunk_idx * activation_chunk_bytes) + chunk_bytes_this_round = min( + activation_chunk_bytes, output_token_bytes - chunk_idx * activation_chunk_bytes + ) + current_chunk_bytes = Int32(chunk_bytes_this_round) + current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) + stateless_data_pacing_enabled = ( + self.token_back_mode != "reuse_dispatch_warps" # nosec B105 + and current_window_unit > self.minimum_pacing_window_cycles + ) + round_start_clock = Int64(0) + if cutlass.const_expr(reuse_data_pacing_enabled or stateless_data_pacing_enabled): + if is_remote_token: + round_start_clock = read_clock64() + else: + round_start_clock = round_start_clock + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(worker_mbarrier, current_chunk_bytes) + tma_load_1d( + worker_activation_stage.iterator, + local_activation_address + chunk_byte_offset, + worker_mbarrier, + current_chunk_bytes, + ) + cute.arch.mbarrier_wait(worker_mbarrier, transfer_phase) + destination_chunk = cute.make_ptr( + cutlass.Uint8, + destination_row.iterator.toint() + chunk_byte_offset, + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + if cutlass.const_expr(self.reduce_topk_in_kernel): + cp_reduce_async_bulk_add_bf16_s2g( + destination_chunk, worker_activation_stage.iterator, current_chunk_bytes + ) + else: + cp_async_bulk_s2g( + destination_chunk, worker_activation_stage.iterator, current_chunk_bytes + ) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + transfer_phase = transfer_phase ^ Int32(1) + if cutlass.const_expr(reuse_data_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + data_average_window = self._adaptive_pace( + data_average_window, current_window, data_low_window, data_high_window + ) + elif cutlass.const_expr(stateless_data_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + self._stateless_pace(Int32(current_window_unit), current_window) + iket.range_pop() + + if cutlass.const_expr(self.token_back_push_sf): + iket.range_push("token_back.push_sf") + local_sf_address = fc2_sf_pointer.toint() + Int64(pool_token_idx) * Int64(output_sf_bytes) + remote_sf = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.pre_reduced_activation_sf.dtype, + self._token_comm_args.pre_reduced_activation_sf.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.pre_reduced_activation_sf.iterator.max_alignment), + ), + self._token_comm_args.pre_reduced_activation_sf.layout, + ) + destination_sf_row = remote_sf[Int64(metadata.src_token), destination_topk, None] + for chunk_idx in cutlass.range_constexpr(sf_chunk_count): + chunk_byte_offset = Int64(chunk_idx * sf_chunk_bytes) + chunk_bytes_this_round = min(sf_chunk_bytes, output_sf_bytes - chunk_idx * sf_chunk_bytes) + current_chunk_bytes = Int32(chunk_bytes_this_round) + current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) + stateless_sf_pacing_enabled = ( + self.token_back_mode != "reuse_dispatch_warps" # nosec B105 + and current_window_unit > self.minimum_pacing_window_cycles + ) + round_start_clock = Int64(0) + if cutlass.const_expr(reuse_sf_pacing_enabled or stateless_sf_pacing_enabled): + if is_remote_token: + round_start_clock = read_clock64() + else: + round_start_clock = round_start_clock + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(worker_mbarrier, current_chunk_bytes) + tma_load_1d( + worker_sf_stage.iterator, + local_sf_address + chunk_byte_offset, + worker_mbarrier, + current_chunk_bytes, + ) + cute.arch.mbarrier_wait(worker_mbarrier, transfer_phase) + destination_chunk = cute.make_ptr( + cutlass.Uint8, + destination_sf_row.iterator.toint() + chunk_byte_offset, + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_chunk, worker_sf_stage.iterator, current_chunk_bytes) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + transfer_phase = transfer_phase ^ Int32(1) + if cutlass.const_expr(reuse_sf_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + sf_average_window = self._adaptive_pace( + sf_average_window, current_window, sf_low_window, sf_high_window + ) + elif cutlass.const_expr(stateless_sf_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + self._stateless_pace(Int32(current_window_unit), current_window) + iket.range_pop() + + cute.arch.sync_warp() + next_dense_token = self.next_token(next_dense_token) + + expert_valid_begin = expert_valid_end + local_expert = local_expert + Int32(1) + iket.range_pop() + + @cute.jit + def next_token(self, current_token: Int32) -> Int32: + global_worker_count = self.promised_launchable_sm_count * self.transfer_warp_count + schedule_counter = None + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 + schedule_counter = self._device_workspace.ptr(self.token_back_schedule_region) + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 + claimed_token = Int32(0) + if self._lane_idx == Int32(0): + claimed_token = cute.arch.atomic_add(schedule_counter, Int32(1), sem="relaxed", scope="gpu") + return Int32(cute.arch.shuffle_sync(claimed_token, Int32(0))) + return current_token + global_worker_count + + @cute.jit + def reset_tail(self) -> None: + """Reset communication state with all configured token-in transfer warps.""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + transfer_thread_idx = transfer_warp_idx * Int32(32) + lane_idx + iket.range_push("tail.nvlink_drain") + self._nvlink_barrier.arrive_and_wait( + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + prologue_grid_sync=True, + epilogue_grid_sync=False, + ) + iket.range_pop() + total_reset_threads = self.promised_launchable_sm_count * self.transfer_thread_count + global_reset_thread = self._linear_cta_idx * Int32(self.transfer_thread_count) + transfer_thread_idx + iket.range_push("tail.reset_workspace") + self._device_workspace.reset_tail_space("shared", global_reset_thread, total_reset_threads) + self._device_workspace.reset_tail_space("local", global_reset_thread, total_reset_threads) + iket.range_pop() + iket.range_push("tail.nvlink_publish") + self._nvlink_barrier.arrive_and_wait( + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + prologue_grid_sync=False, + epilogue_grid_sync=True, + ) + iket.range_pop() + if cutlass.const_expr(os.environ.get("MEGA_USE_NCU", "0") == "1"): + iket.range_push("tail.ncu_finalize") + self._nvlink_barrier.finalize( + 2, + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + ) + iket.range_pop() + + +__all__ = ["TokenCommDeterministic"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/token_protocol.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/token_protocol.py index 8ed3b33dd75d..289c5fe47298 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/token_protocol.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/token_protocol.py @@ -9,7 +9,6 @@ from cutlass.cute.typing import AddressSpace from cutlass.cutlass_dsl import Int32, Int64 - @dataclasses.dataclass(frozen=True) class TokenSrcMetadata: """One i64 routing record: rank:u16, topk:u16, token:u32.""" @@ -22,11 +21,15 @@ class TokenSrcMetadata: def pack(self) -> Int64: high = (Int64(self.src_rank) << Int64(16)) | Int64(self.src_topk) - return (high << Int64(32)) | (Int64(self.src_token) & Int64(0xFFFFFFFF)) + return (high << Int64(32)) | ( + Int64(self.src_token) & Int64(0xFFFFFFFF) + ) @staticmethod def _pointer(address: Union[cute.Pointer, Int64]) -> cute.Pointer: - raw_address = address if isinstance(address, Int64) else address.toint() + raw_address = ( + address if isinstance(address, Int64) else address.toint() + ) return cute.make_ptr( Int64, raw_address, @@ -42,7 +45,9 @@ def load( cls, address: Union[cute.Pointer, Int64], ) -> "TokenSrcMetadata": - packed = Int64(cute.arch.load(cls._pointer(address), Int64, scope="gpu")) + packed = Int64( + cute.arch.load(cls._pointer(address), Int64, scope="gpu") + ) high = packed >> Int64(32) return cls( src_rank=Int32((high >> Int64(16)) & Int64(0xFFFF)), diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/__init__.py index 52a7a9daf028..a84af22de952 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/__init__.py @@ -1,2 +1,22 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Blackwell kernel implementations.""" + +from .inference import ( + BlockScaledSwapAbFc12Extension, + BlockScaledSwapAbFc12Kernel, + BlockScaledSwapAbFc12Mainloop, + BlockScaledSwapAbMegaMoeKernel, + GatedActEpilogueArgs, + SwapABGatedActEpilogue, + TensorRole, +) + + +__all__ = [ + "BlockScaledSwapAbFc12Extension", + "BlockScaledSwapAbFc12Kernel", + "BlockScaledSwapAbFc12Mainloop", + "BlockScaledSwapAbMegaMoeKernel", + "GatedActEpilogueArgs", + "SwapABGatedActEpilogue", + "TensorRole", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/constants.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/constants.py index ce2c01076f54..7870293a0bf7 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/constants.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/constants.py @@ -1,8 +1,6 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Specification-defined numeric constants shared across kernel components.""" + Log2E = 1.4426950408889634 Fp32Max = 3.40282346638528859812e38 diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/cute_py_helpers.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/cute_py_helpers.py index 625075c71335..8cf623735b84 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/cute_py_helpers.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/cute_py_helpers.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - from dataclasses import dataclass from math import gcd from typing import Optional, Tuple, Union @@ -11,6 +8,7 @@ from .smem_workspace import SmemRegion + # Every block-scaled tcgen05 MMA kind accumulates in F32 and no other accumulator is legal, so # this belongs to the instruction family rather than being a caller choice. tcgen05_block_scaled_acc_dtype = cutlass.Float32 @@ -54,9 +52,7 @@ def make_tcgen05_tmem_plan( accumulators; other supported cases maximize disjoint accumulator stages. """ if mma_instruction.participates not in (1, 2): - raise ValueError( - f"TCGen05 MMA participates must be one or two, got {mma_instruction.participates}." - ) + raise ValueError(f"TCGen05 MMA participates must be one or two, got {mma_instruction.participates}.") if len(mma_instruction.instruction_mnk) != 3 or len(mma_tiler_mnk) != 3: raise ValueError("instruction_mnk and mma_tiler_mnk must each contain three dimensions.") if any(dimension <= 0 for dimension in (*mma_instruction.instruction_mnk, *mma_tiler_mnk)): @@ -64,10 +60,7 @@ def make_tcgen05_tmem_plan( instruction_m, instruction_n, instruction_k = mma_instruction.instruction_mnk tile_m, tile_n, tile_k = mma_tiler_mnk - if ( - instruction_m % mma_instruction.participates != 0 - or instruction_n % mma_instruction.participates != 0 - ): + if instruction_m % mma_instruction.participates != 0 or instruction_n % mma_instruction.participates != 0: raise ValueError("MMA instruction M and N must be divisible by participates.") if tile_m % instruction_m != 0 or tile_n % instruction_n != 0 or tile_k % instruction_k != 0: raise ValueError("MMA instruction dimensions must divide mma_tiler_mnk.") @@ -97,9 +90,7 @@ def make_tcgen05_tmem_plan( tmem_column_capacity = cute.arch.get_max_tmem_alloc_cols(arch) accumulator_stage_columns = tile_n scale_factor_columns = sfa_columns + sfb_columns - accumulator_stage_count = ( - tmem_column_capacity - scale_factor_columns - ) // accumulator_stage_columns + accumulator_stage_count = (tmem_column_capacity - scale_factor_columns) // accumulator_stage_columns if accumulator_stage_count < 1: raise ValueError("Scale-factor TMEM leaves no accumulator stage.") @@ -161,9 +152,7 @@ def _round_tmem_allocation_columns(used_columns: int, capacity_columns: int) -> else: allocation_columns = _round_up(used_columns, 32) if allocation_columns > capacity_columns: - raise ValueError( - f"TMEM plan needs {allocation_columns} columns, exceeding {capacity_columns}." - ) + raise ValueError(f"TMEM plan needs {allocation_columns} columns, exceeding {capacity_columns}.") return allocation_columns @@ -192,9 +181,7 @@ def make_smem_layouts( arch: str, ) -> Union[Tuple[SmemRegion, SmemRegion], Tuple[SmemRegion, SmemRegion, SmemRegion, SmemRegion]]: """Derive workspace-ready TCGen05 operand regions without an MLIR context.""" - has_scale_factors = any( - value is not None for value in (mma_inst.sfa_type, mma_inst.sfb_type, mma_inst.sf_vec_size) - ) + has_scale_factors = any(value is not None for value in (mma_inst.sfa_type, mma_inst.sfb_type, mma_inst.sf_vec_size)) if has_scale_factors and any( value is None for value in (mma_inst.sfa_type, mma_inst.sfb_type, mma_inst.sf_vec_size) ): @@ -211,9 +198,7 @@ def make_smem_layouts( raise ValueError(f"All stage counts must be positive, got {operand_stages}.") if mma_inst.participates not in (1, 2): - raise ValueError( - f"TCGen05 MMA participates must be one or two, got {mma_inst.participates}." - ) + raise ValueError(f"TCGen05 MMA participates must be one or two, got {mma_inst.participates}.") if len(mma_inst.instruction_mnk) != 3 or len(mma_tiler_mnk) != 3: raise ValueError("instruction_mnk and mma_tiler_mnk must each contain three dimensions.") if any(dimension <= 0 for dimension in (*mma_inst.instruction_mnk, *mma_tiler_mnk)): @@ -401,12 +386,8 @@ def _tiled_divide_2d(shape, stride, tile_shape): ): raise ValueError("A 2-D tiled divide requires rank-2 shape, stride, and tile.") - mn_tile_stride, mn_rest_shape, mn_rest_stride = _divide_mode_by_tile( - shape[0], stride[0], tile_shape[0] - ) - k_tile_stride, k_rest_shape, k_rest_stride = _divide_mode_by_tile( - shape[1], stride[1], tile_shape[1] - ) + mn_tile_stride, mn_rest_shape, mn_rest_stride = _divide_mode_by_tile(shape[0], stride[0], tile_shape[0]) + k_tile_stride, k_rest_shape, k_rest_stride = _divide_mode_by_tile(shape[1], stride[1], tile_shape[1]) return ( ((tile_shape[0], tile_shape[1]), mn_rest_shape, k_rest_shape), ((mn_tile_stride, k_tile_stride), mn_rest_stride, k_rest_stride), @@ -471,9 +452,7 @@ def _make_ab_region( mma_mn_shape = _prefix_tile_profile(mn_shape, value_shape[0]) mma_k_shape = _prefix_tile_profile(k_shape, value_shape[1]) - mma_shape, mma_stride = _tiled_divide_2d( - (mn_shape, k_shape), (mn_stride, k_stride), (mma_mn_shape, mma_k_shape) - ) + mma_shape, mma_stride = _tiled_divide_2d((mn_shape, k_shape), (mn_stride, k_stride), (mma_mn_shape, mma_k_shape)) return SmemRegion( name="", kind="tensor", @@ -512,10 +491,7 @@ def _make_sf_region( full_k_stride = ((0, 1), 512 if basic_chunk_repetitions > 1 else 0) mma_mn_shape = ((32, 4), instruction_mn_blocks) - mma_k_shape = ( - sf_vec_size, - _prefix_tile_profile((4, basic_chunk_repetitions), instruction_k // sf_vec_size), - ) + mma_k_shape = (sf_vec_size, _prefix_tile_profile((4, basic_chunk_repetitions), instruction_k // sf_vec_size)) mma_shape, mma_stride = _tiled_divide_2d( (full_mn_shape, full_k_shape), (full_mn_stride, full_k_stride), (mma_mn_shape, mma_k_shape) ) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/device_workspace.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/device_workspace.py index ec982a384ac9..df72be3e4037 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/device_workspace.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/device_workspace.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Declarative GMEM workspace shared by host layout and device access.""" import dataclasses @@ -18,11 +15,12 @@ is_nested_shape, is_power_of_two, ordered_stride, - round_up, row_major_stride, + round_up, validate_static_integer_tuple, ) + BufferResetAttr = Literal["data", "zero_on_first_allocate", "tail_reset"] BufferSpace = Literal["local", "shared"] @@ -144,18 +142,14 @@ def finalize(self) -> None: if self._finalized: raise RuntimeError("DeviceWorkspace.finalize() may only be called once.") for buffer_space in ("local", "shared"): - ordered = sorted( - self._registered[buffer_space], key=lambda region: _reset_order[region.reset] - ) + ordered = sorted(self._registered[buffer_space], key=lambda region: _reset_order[region.reset]) cursor = 0 for region_index, region in enumerate(ordered): stride = region.stride if stride is None: raise RuntimeError(f"Region {region.name!r} stride was not resolved.") cosize = cosize_from_shape_stride_tuples(region.shape, stride) - nbytes = ( - _footprint_from_shape_stride(region.shape, stride) * int(region.dtype.width) + 7 - ) // 8 + nbytes = (_footprint_from_shape_stride(region.shape, stride) * int(region.dtype.width) + 7) // 8 if region.reset == "tail_reset" and ( region_index == 0 or ordered[region_index - 1].reset != "tail_reset" ): @@ -170,8 +164,7 @@ def finalize(self) -> None: self._byte_alignment[region.name] = byte_alignment cursor += nbytes is_last_tail_reset_region = region.reset == "tail_reset" and ( - region_index + 1 == len(ordered) - or ordered[region_index + 1].reset != "tail_reset" + region_index + 1 == len(ordered) or ordered[region_index + 1].reset != "tail_reset" ) if is_last_tail_reset_region: cursor = round_up(cursor, 16) @@ -246,16 +239,12 @@ def ptr(self, name: str) -> cute.Pointer: region = self._region_by_name[name] base = self._base[region.buffer_space] address = base.toint() + Int64(self._offset[name]) - return cute.make_ptr( - region.dtype, address, AddressSpace.gmem, assumed_align=self._byte_alignment[name] - ) + return cute.make_ptr(region.dtype, address, AddressSpace.gmem, assumed_align=self._byte_alignment[name]) @cute.jit def tensor(self, name: str) -> cute.Tensor: region = self._region_by_name[name] - return cute.make_tensor( - self.ptr(name), cute.make_layout(region.shape, stride=self._stride[name]) - ) + return cute.make_tensor(self.ptr(name), cute.make_layout(region.shape, stride=self._stride[name])) @cute.jit def reset_tail(self, tid: Int32, total_threads: int) -> None: @@ -267,17 +256,13 @@ def reset_tail_space(self, buffer_space: BufferSpace, tid: Int32, total_threads: num_vectors = self._tail_leading[buffer_space] // 16 if cutlass.const_expr(num_vectors > 0): vectors = cute.make_tensor( - cute.make_ptr( - Int32, self._base[buffer_space].toint(), AddressSpace.gmem, assumed_align=16 - ), + cute.make_ptr(Int32, self._base[buffer_space].toint(), AddressSpace.gmem, assumed_align=16), cute.make_layout((num_vectors, 4), stride=(4, 1)), ) zero = cute.make_rmem_tensor((4,), Int32) for element in cutlass.range_constexpr(4): zero[element] = Int32(0) - store_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), Int32, num_bits_per_copy=128 - ) + store_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), Int32, num_bits_per_copy=128) reset_round_count = ceil_div(num_vectors, total_threads) for reset_round in cutlass.range_constexpr(reset_round_count): vector_index = Int32(reset_round * total_threads) + tid diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/dlb_expert_count.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/dlb_expert_count.py new file mode 100644 index 000000000000..596362cf6bb9 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/dlb_expert_count.py @@ -0,0 +1,30 @@ +"""Resolve logical experts and helper capacity at the host construction boundary.""" + +from ..api import ImplDesc, ProblemDesc + + +def resolve_dlb_expert_counts( + problem_desc: ProblemDesc, impl_desc: ImplDesc, *, world_size: int +) -> tuple[ProblemDesc, ImplDesc]: + """Normalize the new logical-count API while retaining legacy physical counts.""" + if "total_helper_slots" not in impl_desc: + return problem_desc, impl_desc + if "helper_expert_count" in impl_desc: + raise ValueError("Use total_helper_slots or legacy helper_expert_count, not both.") + + logical_experts = problem_desc["expert_count"] + total_helpers = impl_desc["total_helper_slots"] + if type(world_size) is not int or world_size <= 0: + raise ValueError("world_size must be a positive exact int.") + if type(logical_experts) is not int or logical_experts <= 0: + raise ValueError("Logical expert_count must be a positive exact int.") + if type(total_helpers) is not int or total_helpers < 0: + raise ValueError("total_helper_slots must be a non-negative exact int.") + if logical_experts % world_size or total_helpers % world_size: + raise ValueError("Logical expert_count and total_helper_slots must be divisible by world_size.") + + physical_problem = ProblemDesc({**problem_desc, "expert_count": logical_experts + total_helpers}) + physical_impl = dict(impl_desc) + del physical_impl["total_helper_slots"] + physical_impl["helper_expert_count"] = total_helpers // world_size + return physical_problem, ImplDesc(physical_impl) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/dsl_helpers.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/dsl_helpers.py index bdb83e4b0641..807d85c27b59 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/dsl_helpers.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/dsl_helpers.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """General-purpose CuTe DSL helpers.""" from typing import Callable, Literal, Optional @@ -28,12 +25,8 @@ def smem_exclusive_prefix( num_elements = cute.size(input_tensor) scan_rows = num_elements // 4 num_warps = block_thread_count // 32 - input_vectors = cute.make_tensor( - input_tensor.iterator, cute.make_layout((scan_rows, 4), stride=(4, 1)) - ) - load_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=128 - ) + input_vectors = cute.make_tensor(input_tensor.iterator, cute.make_layout((scan_rows, 4), stride=(4, 1))) + load_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=128) values = cute.make_rmem_tensor((4,), cutlass.Int32) carry = Int32(0) @@ -51,12 +44,7 @@ def smem_exclusive_prefix( for element in cutlass.range_constexpr(4): values[element] = Int32(0) - local_prefix = ( - Int32(0), - values[0], - values[0] + values[1], - values[0] + values[1] + values[2], - ) + local_prefix = (Int32(0), values[0], values[0] + values[1], values[0] + values[1] + values[2]) lane_total = local_prefix[3] + values[3] inclusive = lane_total for step_log in cutlass.range_constexpr(5): @@ -95,17 +83,12 @@ def smem_exclusive_prefix( def mark_alignment(tensor: cute.Tensor, byte_alignment: int) -> cute.Tensor: pointer = tensor.iterator return cute.make_tensor( - cute.make_ptr( - pointer.dtype, pointer.toint(), pointer.memspace, assumed_align=byte_alignment - ), - tensor.layout, + cute.make_ptr(pointer.dtype, pointer.toint(), pointer.memspace, assumed_align=byte_alignment), tensor.layout ) @cute.jit -def spin_peek( - pointer: Pointer, condition: Callable[[Int32], Boolean], scope: str = "gpu" -) -> Boolean: +def spin_peek(pointer: Pointer, condition: Callable[[Int32], Boolean], scope: str = "gpu") -> Boolean: """Perform one acquire load and test its value.""" value = cute.arch.load(pointer, pointer.dtype, sem="acquire", scope=scope) return Boolean(condition(value)) @@ -138,10 +121,7 @@ def _tma_multicast_pattern( if tensor_role in ("a", "sfa"): return sum(1 << (cluster_n_index * cluster_m) for cluster_n_index in range(cluster_n)) if tensor_role == "b": - return sum( - 1 << (cluster_m_index * mma_cta_count) - for cluster_m_index in range(cluster_m // mma_cta_count) - ) + return sum(1 << (cluster_m_index * mma_cta_count) for cluster_m_index in range(cluster_m // mma_cta_count)) if tensor_role == "sfb": return (1 << cluster_m) - 1 raise ValueError(f"Unsupported TMA tensor role {tensor_role!r}.") @@ -176,11 +156,7 @@ def tma_multicast_mask( ) else: cta_n = Int32(cta_coord_in_cluster[1]) - mma_cta_index = ( - Int32(0) - if cutlass.const_expr(tensor_role == "sfb" or mma_cta_count == 1) - else cta_m % 2 - ) + mma_cta_index = Int32(0) if cutlass.const_expr(tensor_role == "sfb" or mma_cta_count == 1) else cta_m % 2 result = cute.arch.inline_ptx( "{\n\t" ".reg .u32 offset;\n\t" @@ -210,11 +186,7 @@ def tma_multicast_mask( ) else: cta_n = Int32(cta_coord_in_cluster[1]) - mma_cta_index = ( - Int32(0) - if cutlass.const_expr(tensor_role == "sfb" or mma_cta_count == 1) - else cta_m % 2 - ) + mma_cta_index = Int32(0) if cutlass.const_expr(tensor_role == "sfb" or mma_cta_count == 1) else cta_m % 2 result = cute.arch.inline_ptx( "{\n\t" ".reg .u32 offset;\n\t" @@ -227,9 +199,7 @@ def tma_multicast_mask( return Int16(result) if cutlass.const_expr(is_preferred is None): - raise ValueError( - "is_preferred is required when the preferred and fallback multicast patterns differ." - ) + raise ValueError("is_preferred is required when the preferred and fallback multicast patterns differ.") if cutlass.const_expr(tensor_role in ("a", "sfa")): result = cute.arch.inline_ptx( @@ -243,15 +213,9 @@ def tma_multicast_mask( ) else: if cutlass.const_expr(fallback_pattern & preferred_pattern != fallback_pattern): - raise ValueError( - "Fallback B/SFB multicast pattern must be a subset of the preferred pattern." - ) + raise ValueError("Fallback B/SFB multicast pattern must be a subset of the preferred pattern.") cta_n = Int32(cta_coord_in_cluster[1]) - mma_cta_index = ( - Int32(0) - if cutlass.const_expr(tensor_role == "sfb" or mma_cta_count == 1) - else cta_m % 2 - ) + mma_cta_index = Int32(0) if cutlass.const_expr(tensor_role == "sfb" or mma_cta_count == 1) else cta_m % 2 result = cute.arch.inline_ptx( "{\n\t" ".reg .u32 cluster_m, offset;\n\t" @@ -268,10 +232,4 @@ def tma_multicast_mask( return Int16(result) -__all__ = [ - "mark_alignment", - "smem_exclusive_prefix", - "spin_peek", - "spin_wait", - "tma_multicast_mask", -] +__all__ = ["mark_alignment", "smem_exclusive_prefix", "spin_peek", "spin_wait", "tma_multicast_mask"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/fc2_ready_geometry.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/fc2_ready_geometry.py new file mode 100644 index 000000000000..0ec5abde4407 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/fc2_ready_geometry.py @@ -0,0 +1,36 @@ +"""Host geometry for FC2 completion counters in the swap-AB kernel.""" + + +def derive_fc2_ready_geometry( + *, + hidden_size: int, + mma_tiler_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + use_2cta_instrs: bool, + fallback_cluster_shape_mn: tuple[int, int] | None = None, +) -> tuple[int, int]: + """Return (tokens per ready slot, CTA completions per ready slot). + + Work IDs and token slots use the preferred *logical* cluster. Fixed groups + of fallback clusters partition that cluster's CTA coordinates; they do not + create extra token slots or smaller completion targets. Every logical CTA + publishes once, including padded hidden tiles in the last cluster tile. + """ + mma_cta_count = 2 if use_2cta_instrs else 1 + if cluster_shape_mn[1] != 1: + raise ValueError("FC2 token-ready geometry requires hidden-only cluster splitting (N=1).") + if mma_tiler_mnk[0] % mma_cta_count or cluster_shape_mn[0] % mma_cta_count: + raise ValueError("FC2 token-ready geometry must preserve complete MMA CTA groups.") + if fallback_cluster_shape_mn is not None: + fallback_m, fallback_n = fallback_cluster_shape_mn + if ( + fallback_n != 1 + or fallback_m <= 0 + or cluster_shape_mn[0] % fallback_m + or fallback_m % mma_cta_count + ): + raise ValueError("FC2 fallback clusters must partition the preferred hidden axis into complete MMA groups.") + cta_hidden = mma_tiler_mnk[0] // mma_cta_count + hidden_per_logical_cluster = cta_hidden * cluster_shape_mn[0] + hidden_cluster_tiles = (hidden_size + hidden_per_logical_cluster - 1) // hidden_per_logical_cluster + return mma_tiler_mnk[1], hidden_cluster_tiles * cluster_shape_mn[0] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/flag_batch.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/flag_batch.py index b38730b01cfc..28088ec8f0ca 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/flag_batch.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/flag_batch.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Lane- and warp-distributed GPU release-counter batching.""" import dataclasses @@ -24,22 +21,15 @@ class GpuReleaseFlagBatchTracker: thread_idx: Int32 @cute.jit - def _make( - self, flag_address: Int64, accumulated_flags: Int32, phase: Int32 - ) -> "GpuReleaseFlagBatchTracker": + def _make(self, flag_address: Int64, accumulated_flags: Int32, phase: Int32) -> "GpuReleaseFlagBatchTracker": return type(self)( - flag_address=flag_address, - accumulated_flags=accumulated_flags, - phase=phase, - thread_idx=self.thread_idx, + flag_address=flag_address, accumulated_flags=accumulated_flags, phase=phase, thread_idx=self.thread_idx ) @cute.jit def fire(self) -> None: if self.flag_address != Int64(0): - pointer = cute.make_ptr( - cutlass.Int32, self.flag_address, AddressSpace.gmem, assumed_align=4 - ) + pointer = cute.make_ptr(cutlass.Int32, self.flag_address, AddressSpace.gmem, assumed_align=4) red_add_release_gpu_s32(pointer, Int32(1)) @cute.jit @@ -51,12 +41,8 @@ def accumulate( lane_address = Int64(0) if self.thread_idx == Int32(0): lane_address = flag_address - self._make( - flag_address=lane_address, accumulated_flags=Int32(1), phase=self.phase - ).fire() - return self._make( - flag_address=Int64(0), accumulated_flags=Int32(0), phase=Int32(next_phase) - ) + self._make(flag_address=lane_address, accumulated_flags=Int32(1), phase=self.phase).fire() + return self._make(flag_address=Int64(0), accumulated_flags=Int32(0), phase=Int32(next_phase)) current_address = self.flag_address accumulated_flags = self.accumulated_flags @@ -66,19 +52,11 @@ def accumulate( if accumulated_flags == Int32(flush_threshold) or next_phase != self.phase: if cutlass.const_expr(not no_fire): - self._make( - flag_address=current_address, - accumulated_flags=accumulated_flags, - phase=self.phase, - ).fire() + self._make(flag_address=current_address, accumulated_flags=accumulated_flags, phase=self.phase).fire() accumulated_flags = Int32(0) current_address = Int64(0) - return self._make( - flag_address=current_address, - accumulated_flags=accumulated_flags, - phase=Int32(next_phase), - ) + return self._make(flag_address=current_address, accumulated_flags=accumulated_flags, phase=Int32(next_phase)) @dataclasses.dataclass(frozen=True) @@ -91,22 +69,15 @@ class GpuAsyncReleaseFlagBatchTracker: warp_idx: Int32 @cute.jit - def _make( - self, flag_address: Int64, accumulated_flags: Int32, phase: Int32 - ) -> "GpuAsyncReleaseFlagBatchTracker": + def _make(self, flag_address: Int64, accumulated_flags: Int32, phase: Int32) -> "GpuAsyncReleaseFlagBatchTracker": return type(self)( - flag_address=flag_address, - accumulated_flags=accumulated_flags, - phase=phase, - warp_idx=self.warp_idx, + flag_address=flag_address, accumulated_flags=accumulated_flags, phase=phase, warp_idx=self.warp_idx ) @cute.jit def fire(self) -> None: if self.flag_address != Int64(0): - pointer = cute.make_ptr( - cutlass.Int32, self.flag_address, AddressSpace.gmem, assumed_align=4 - ) + pointer = cute.make_ptr(cutlass.Int32, self.flag_address, AddressSpace.gmem, assumed_align=4) with cute.arch.elect_one(): red_async_add_release_gpu_s32(pointer, Int32(1)) @@ -119,12 +90,8 @@ def accumulate( warp_address = Int64(0) if self.warp_idx == Int32(0): warp_address = flag_address - self._make( - flag_address=warp_address, accumulated_flags=Int32(1), phase=self.phase - ).fire() - return self._make( - flag_address=Int64(0), accumulated_flags=Int32(0), phase=Int32(next_phase) - ) + self._make(flag_address=warp_address, accumulated_flags=Int32(1), phase=self.phase).fire() + return self._make(flag_address=Int64(0), accumulated_flags=Int32(0), phase=Int32(next_phase)) current_address = self.flag_address accumulated_flags = self.accumulated_flags @@ -134,29 +101,16 @@ def accumulate( if accumulated_flags == Int32(flush_threshold) or next_phase != self.phase: if cutlass.const_expr(not no_fire): - self._make( - flag_address=current_address, - accumulated_flags=accumulated_flags, - phase=self.phase, - ).fire() + self._make(flag_address=current_address, accumulated_flags=accumulated_flags, phase=self.phase).fire() accumulated_flags = Int32(0) current_address = Int64(0) - return self._make( - flag_address=current_address, - accumulated_flags=accumulated_flags, - phase=Int32(next_phase), - ) + return self._make(flag_address=current_address, accumulated_flags=accumulated_flags, phase=Int32(next_phase)) @cute.jit def make_flag_batch_tracker( - use_async: bool, - *, - flag_address: Int64, - accumulated_flags: Int32, - phase: Int32, - thread_idx: Int32, + use_async: bool, *, flag_address: Int64, accumulated_flags: Int32, phase: Int32, thread_idx: Int32 ) -> Union[GpuReleaseFlagBatchTracker, GpuAsyncReleaseFlagBatchTracker]: """Construct a lane-batched synchronous or warp-batched asynchronous tracker. @@ -172,15 +126,8 @@ def make_flag_batch_tracker( warp_idx=cute.arch.make_warp_uniform(thread_idx // Int32(32)), ) return GpuReleaseFlagBatchTracker( - flag_address=flag_address, - accumulated_flags=accumulated_flags, - phase=phase, - thread_idx=thread_idx, + flag_address=flag_address, accumulated_flags=accumulated_flags, phase=phase, thread_idx=thread_idx ) -__all__ = [ - "GpuAsyncReleaseFlagBatchTracker", - "GpuReleaseFlagBatchTracker", - "make_flag_batch_tracker", -] +__all__ = ["GpuAsyncReleaseFlagBatchTracker", "GpuReleaseFlagBatchTracker", "make_flag_batch_tracker"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/helper_weight_ready_gate.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/helper_weight_ready_gate.py new file mode 100644 index 000000000000..25122eb351e2 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/helper_weight_ready_gate.py @@ -0,0 +1,318 @@ +"""Shared Blackwell/Rubin kernel gate for dynamically installed helper weights. + +The gate is deliberately optional at code generation time. A specialization +with ``S == 0`` constructs no gate and therefore carries neither of the two +runtime values. A helper-bearing specialization derives ``H = M - S`` and +waits once per scheduler CTA, immediately before its first physical helper +work tile is published. + +``terminal_flags`` is the local ``uint64[EP]`` view published by the weight +transport. Bit 63 is reserved. A source is terminal-ready when that bit is +clear and its low-63-bit generation is greater than or equal to the expected +generation. The expected generation is a CPU-captured launch-time scalar, not +device READY storage. After READY publication, callers keep the weights +immutable until the MegaMoE launch completes. +""" + +import dataclasses +from enum import IntEnum +from typing import Optional + +import cutlass +import cutlass.cute as cute +from cutlass._mlir.dialects import llvm +from cutlass.cute.typing import AddressSpace, Pointer +from cutlass.cutlass_dsl import Int32, extract_mlir_values, new_from_mlir_values + +from .ptx_helpers import nanosleep, read_globaltimer_ns + + +HELPER_WEIGHT_READY_TERMINAL_ERROR_BIT = 1 << 63 +HELPER_WEIGHT_READY_GENERATION_MASK = HELPER_WEIGHT_READY_TERMINAL_ERROR_BIT - 1 +HELPER_WEIGHT_READY_POLL_SLEEP_NS = 100 +# Elapsed time after the first unsuccessful READY sweep, independent of SM clocks. +# This is deliberately compile-time policy, not another production ABI knob. +HELPER_WEIGHT_READY_TIMEOUT_NS = 120_000_000_000 + + +class HelperWeightReadyStatus(IntEnum): + """Warp-uniform result of one terminal-flags sweep.""" + + WAITING = 0 + READY = 1 + TERMINAL_ERROR = 2 + INVALID_EXPECTED_GENERATION = 3 + TIMEOUT = 4 + + +@dataclasses.dataclass(frozen=True) +class HelperWeightReadyTopology: + """Compile-time ``[H home slots][S helper slots]`` topology.""" + + memory_slot_count: int + helper_count: int + home_expert_count: int + gate_required: bool + +def derive_helper_weight_ready_topology( + memory_slot_count: int, + helper_count: int, +) -> HelperWeightReadyTopology: + """Validate static ``M``/``S`` and derive the only helper boundary ``H``.""" + + if type(memory_slot_count) is not int or memory_slot_count <= 0: + raise ValueError("memory_slot_count (M) must be a positive exact int") + if type(helper_count) is not int or helper_count < 0: + raise ValueError("helper_count (S) must be a non-negative exact int") + if helper_count >= memory_slot_count: + raise ValueError("helper_count (S) must leave at least one home slot") + + home_expert_count = memory_slot_count - helper_count + return HelperWeightReadyTopology( + memory_slot_count=memory_slot_count, + helper_count=helper_count, + home_expert_count=home_expert_count, + gate_required=helper_count > 0, + ) + + +def _validate_helper_weight_ready_binding( + *, + topology: HelperWeightReadyTopology, + source_count: Optional[int], + terminal_flags: Optional[Pointer], + expected_generation: Optional[object], +) -> None: + """Enforce exact runtime-bundle presence from the compile-time topology.""" + + if not topology.gate_required: + if terminal_flags is not None or expected_generation is not None: + raise ValueError("S=0 must not carry READY runtime ABI values") + return + + if type(source_count) is not int or source_count <= 0: + raise ValueError("source_count (EP) must be a positive exact int when S>0") + if terminal_flags is None or expected_generation is None: + raise ValueError("S>0 requires terminal_flags and expected_generation") + if isinstance(expected_generation, int): + if type(expected_generation) is not int or not ( + 1 <= expected_generation <= HELPER_WEIGHT_READY_GENERATION_MASK + ): + raise ValueError("expected_generation must be in [1, 2**63)") + + +@cute.jit +def _helper_weight_ready_fail_closed_trap() -> None: + """Prevent publication after a terminal, invalid, or timed-out wait.""" + + llvm.inline_asm( + res=None, + operands_=[], + asm_string="trap;", + constraints="", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +class HelperWeightReadyGenerationGate: + """One scheduler CTA's wait-all gate with an explicit caller-owned READY latch.""" + + def __init__( + self, + *, + topology: HelperWeightReadyTopology, + source_count: int, + terminal_flags: Pointer, + expected_generation, + ) -> None: + _validate_helper_weight_ready_binding( + topology=topology, + source_count=source_count, + terminal_flags=terminal_flags, + expected_generation=expected_generation, + ) + if not topology.gate_required: + raise ValueError("construct no HelperWeightReadyGenerationGate when S=0") + + self.topology = topology + self.source_count = source_count + self.terminal_flags = terminal_flags + self.expected_generation = expected_generation + + def __extract_mlir_values__(self) -> list: + values = [] + for field in ( + self.terminal_flags, + self.expected_generation, + ): + values.extend(extract_mlir_values(field)) + return values + + def __new_from_mlir_values__(self, values: list) -> "HelperWeightReadyGenerationGate": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + result = new_from_mlir_values( + field, + values[value_index : value_index + field_value_count], + ) + value_index += field_value_count + return result + + result = type(self)( + topology=self.topology, + source_count=self.source_count, + terminal_flags=rebuild(self.terminal_flags), + expected_generation=rebuild(self.expected_generation), + ) + if value_index != len(values): + raise ValueError( + "HelperWeightReadyGenerationGate MLIR value count mismatch: " + f"consumed {value_index}, got {len(values)}." + ) + return result + + @cute.jit + def _typed_terminal_flags(self): + return cute.make_ptr( + cutlass.Uint64, + self.terminal_flags.toint(), + AddressSpace.gmem, + assumed_align=8, + ) + + @cute.jit + def _sweep_terminal_flags(self): + """Read every source cell; acquire all releases before admitting READY.""" + + # The launch generation is a uniform scalar. This terminal pointer is + # the only global-memory READY authority loaded by the spin loop. + flags = self._typed_terminal_flags() + all_sources_geq = Int32(1) + terminal_error_seen = Int32(0) + generation_mask = cutlass.Uint64(HELPER_WEIGHT_READY_GENERATION_MASK) + + for source_rank in cutlass.range_constexpr(self.source_count): + observed = cute.arch.load( + flags + Int32(source_rank), + cutlass.Uint64, + sem="relaxed", + scope="sys", + ) + observed_generation = observed & generation_mask + if observed > generation_mask: + terminal_error_seen = Int32(1) + if observed_generation < self.expected_generation: + all_sources_geq = Int32(0) + + status = Int32(HelperWeightReadyStatus.WAITING) + if terminal_error_seen != Int32(0): + status = Int32(HelperWeightReadyStatus.TERMINAL_ERROR) + elif all_sources_geq != Int32(0): + # Each system-scope strong read plus this acquire fence forms an + # acquire pattern for its source's release (PTX ISA, section 8.8). + llvm.inline_asm( + res=None, + operands_=[], + asm_string="fence.acquire.sys;", + constraints="~{memory}", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + status = Int32(HelperWeightReadyStatus.READY) + return status + + @cute.jit + def _wait_all(self): + """Lane 0 polls; fence the READY acquire before returning success.""" + + status = Int32(HelperWeightReadyStatus.WAITING) + if cute.arch.lane_idx() == Int32(0): + generation_mask = cutlass.Uint64(HELPER_WEIGHT_READY_GENERATION_MASK) + if (self.expected_generation == cutlass.Uint64(0)) | ( + self.expected_generation > generation_mask + ): + status = Int32(HelperWeightReadyStatus.INVALID_EXPECTED_GENERATION) + else: + status = self._sweep_terminal_flags() + if status == Int32(HelperWeightReadyStatus.WAITING): + # Already-ready weights take no timer reads. Unsigned elapsed + # time also handles wraparound of the 64-bit nanosecond timer. + wait_started_ns = read_globaltimer_ns() + while status == Int32(HelperWeightReadyStatus.WAITING): + elapsed_ns = read_globaltimer_ns() - wait_started_ns + if elapsed_ns >= cutlass.Uint64(HELPER_WEIGHT_READY_TIMEOUT_NS): + status = Int32(HelperWeightReadyStatus.TIMEOUT) + else: + nanosleep(HELPER_WEIGHT_READY_POLL_SLEEP_NS) + status = self._sweep_terminal_flags() + + status = Int32(cute.arch.shuffle_sync(status, offset=0)) + if status == Int32(HelperWeightReadyStatus.READY): + if cute.arch.lane_idx() == Int32(0): + cute.arch.fence_proxy("async.global") + return status + + @cute.jit + def wait_before_first_helper_publish(self, expert_idx, ready_latched: Int32) -> Int32: + """Return the next CTA latch after admitting physical ``m in [H,M)``. + + The scheduler must carry this scalar through its work loop. Mutating a + gate nested in a frozen extension does not produce a CuTe loop result. + """ + + needs_admission = ( + (expert_idx >= Int32(self.topology.home_expert_count)) + & (expert_idx < Int32(self.topology.memory_slot_count)) + & (ready_latched == Int32(0)) + ) + if needs_admission: + status = self._wait_all() + if status != Int32(HelperWeightReadyStatus.READY): + _helper_weight_ready_fail_closed_trap() + ready_latched = Int32(1) + return ready_latched + + +def make_helper_weight_ready_generation_gate( + *, + memory_slot_count: int, + helper_count: int, + source_count: Optional[int] = None, + terminal_flags: Optional[Pointer] = None, + expected_generation: Optional[object] = None, +) -> Optional[HelperWeightReadyGenerationGate]: + """Create the mandatory ``S>0`` gate or the ABI-free ``S==0`` sentinel.""" + + topology = derive_helper_weight_ready_topology(memory_slot_count, helper_count) + _validate_helper_weight_ready_binding( + topology=topology, + source_count=source_count, + terminal_flags=terminal_flags, + expected_generation=expected_generation, + ) + if not topology.gate_required: + return None + return HelperWeightReadyGenerationGate( + topology=topology, + source_count=source_count, + terminal_flags=terminal_flags, + expected_generation=expected_generation, + ) + + +__all__ = [ + "HELPER_WEIGHT_READY_GENERATION_MASK", + "HELPER_WEIGHT_READY_TIMEOUT_NS", + "HELPER_WEIGHT_READY_TERMINAL_ERROR_BIT", + "HelperWeightReadyGenerationGate", + "HelperWeightReadyStatus", + "HelperWeightReadyTopology", + "derive_helper_weight_ready_topology", + "make_helper_weight_ready_generation_gate", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/iket_compat.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/iket_compat.py index cfcce57b687a..d4cfd4c972b9 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/iket_compat.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/iket_compat.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Compatibility wrapper for optional in-kernel event tracing support.""" try: @@ -9,7 +6,6 @@ try: from cutlass.cute import iket # type: ignore except (ImportError, NotImplementedError): - class _IketShim: """No-op IKET interface for toolchains without the dialect.""" diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/megamoe_aot.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/megamoe_aot.py new file mode 100644 index 000000000000..a2e46b2ecbf5 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/megamoe_aot.py @@ -0,0 +1,14 @@ +"""Restore argument names and READY defaults for loaded MegaMoE functions.""" + +from inspect import signature + + +def make_megamoe_aot_callable(function, kernel): + """Use TVM-FFI's native adapter without attaching scheduler metadata.""" + from tvm_ffi.utils.kwargs_wrapper import make_kwargs_wrapper + + argument_names = list(signature(kernel.__call__).parameters) + # Home-only entries can omit the unused READY pair, as with direct JIT + # calls. Helper-enabled entries require both READY operands explicitly. + defaults = (None, None) if kernel.helper_expert_count == 0 else () + return make_kwargs_wrapper(function, argument_names, arg_defaults=defaults) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/ptx_helpers.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/ptx_helpers.py index 055d71d7c48d..85d392151038 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/ptx_helpers.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/ptx_helpers.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Minimal inline-PTX primitives required by the greenfield NVFP4 kernels.""" from typing import Optional @@ -11,6 +8,7 @@ from cutlass._mlir.dialects import arith, llvm, vector from cutlass.cutlass_dsl import Float32, Int32, Int64, T, dsl_user_op + TmaCacheHintEvictFirst = 0x12F0000000000000 @@ -21,10 +19,8 @@ def _address_value(pointer_or_address, *, loc=None, ip=None): @dsl_user_op -def nanosleep( - sleep_cycles: int, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None -) -> None: - """Suspend the calling thread for up to the requested clock cycles.""" +def nanosleep(sleep_cycles: int, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None) -> None: + """Suspend the calling thread for approximately the requested nanoseconds.""" if cutlass.const_expr(hasattr(cute.arch, "nanosleep")): cute.arch.nanosleep(sleep_time=sleep_cycles, loc=loc, ip=ip) return @@ -58,20 +54,21 @@ def exit(*, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = @dsl_user_op -def read_clock64( - *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None -) -> Int64: +def read_clock64(*, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None) -> Int64: """Read the per-SM 64-bit cycle counter.""" return Int64( llvm.inline_asm( - T.i64(), - [], - "mov.u64 $0, %clock64;", - "=l", - has_side_effects=True, - asm_dialect=0, - loc=loc, - ip=ip, + T.i64(), [], "mov.u64 $0, %clock64;", "=l", has_side_effects=True, asm_dialect=0, loc=loc, ip=ip + ) + ) + + +@dsl_user_op +def read_globaltimer_ns(*, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None) -> cutlass.Uint64: + """Read the 64-bit global nanosecond timer, independent of SM clock rate.""" + return cutlass.Uint64( + llvm.inline_asm( + T.i64(), [], "mov.u64 $0, %globaltimer;", "=l", has_side_effects=True, asm_dialect=0, loc=loc, ip=ip ) ) @@ -80,9 +77,7 @@ def read_clock64( def movmatrix_b16(input_regs: cute.Tensor) -> cute.Tensor: """Transpose every packed m8n8 b16 register fragment across the warp.""" if cutlass.const_expr(input_regs.element_type.width != 32): - raise TypeError( - f"movmatrix_b16 expects packed 32-bit registers, got {input_regs.element_type}." - ) + raise TypeError(f"movmatrix_b16 expects packed 32-bit registers, got {input_regs.element_type}.") input_words = cute.coalesce(cute.flatten(cute.recast_tensor(input_regs, Int32))) output_words = cute.make_rmem_tensor((cute.size(input_words),), Int32) @@ -145,9 +140,7 @@ def cvt_f32_to_fp8_to_f32( @dsl_user_op -def tma_load_1d( - destination_smem, source_gmem, mbarrier_smem, num_bytes, *, loc=None, ip=None -) -> None: +def tma_load_1d(destination_smem, source_gmem, mbarrier_smem, num_bytes, *, loc=None, ip=None) -> None: """Issue a cache-hinted 1D GMEM-to-SMEM bulk copy.""" llvm.inline_asm( None, @@ -330,9 +323,7 @@ def lds128_v4_b32(smem_pointer, *, loc=None, ip=None): @dsl_user_op -def stg_f32( - address: Int64, value: Float32, predicate: Optional[Int32] = None, *, loc=None, ip=None -) -> None: +def stg_f32(address: Int64, value: Float32, predicate: Optional[Int32] = None, *, loc=None, ip=None) -> None: if predicate is None: llvm.inline_asm( None, @@ -358,9 +349,7 @@ def stg_f32( @dsl_user_op -def stg_b64( - address: Int64, value: Int64, predicate: Optional[Int32] = None, *, loc=None, ip=None -) -> None: +def stg_b64(address: Int64, value: Int64, predicate: Optional[Int32] = None, *, loc=None, ip=None) -> None: if predicate is None: llvm.inline_asm( None, @@ -401,11 +390,7 @@ def red_add_relaxed_sys_s32(address: Int64, value: Int32, *, loc=None, ip=None) @dsl_user_op def red_add_relaxed_sys_f32( - address, - value: Float32, - *, - loc: Optional[ir.Location] = None, - ip: Optional[ir.InsertionPoint] = None, + address, value: Float32, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None ) -> None: llvm.inline_asm( None, @@ -436,11 +421,7 @@ def red_add_release_sys_s32(address: Int64, value: Int32, *, loc=None, ip=None) @dsl_user_op def red_add_release_gpu_s32( - counter_pointer, - value: Int32, - *, - loc: Optional[ir.Location] = None, - ip: Optional[ir.InsertionPoint] = None, + counter_pointer, value: Int32, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None ) -> None: llvm.inline_asm( None, @@ -457,11 +438,7 @@ def red_add_release_gpu_s32( @dsl_user_op def red_async_add_release_gpu_s32( - counter_pointer, - value: Int32, - *, - loc: Optional[ir.Location] = None, - ip: Optional[ir.InsertionPoint] = None, + counter_pointer, value: Int32, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None ) -> None: llvm.inline_asm( None, @@ -503,12 +480,7 @@ def red_async_add_release_sys_u32(address: Int64, value: Int32, *, loc=None, ip= @dsl_user_op def red_add_relaxed_sys_v2_bf16x2( - address, - value0, - value1, - *, - loc: Optional[ir.Location] = None, - ip: Optional[ir.InsertionPoint] = None, + address, value0, value1, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None ) -> None: llvm.inline_asm( None, @@ -561,6 +533,136 @@ def cvt_f32x4_to_f8x4_pack_i32(fp32x4: cute.Tensor, fp8_type, *, loc=None, ip=No return Int32(packed_i32) +@dsl_user_op +def cvt_f32_to_ue8m0_raw_i32( + value, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Int32: + """Round one f32 up to a raw UE8M0 byte in the low bits of an i32.""" + packed_scale = llvm.inline_asm( + T.i32(), + [Float32(value).ir_value(loc=loc, ip=ip)], + "{\n" + " .reg .b16 sf2;\n" + " cvt.rp.satfinite.ue8m0x2.f32 sf2, 0f00000000, $1;\n" + " cvt.u32.u16 $0, sf2;\n" + "}", + "=r,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return Int32(packed_scale) + + +@cute.jit +def cvt_f32x4_to_f8x4_scaled_pack_i32( + fp32x4: cute.Tensor, + fp8_type, + scale_raw_i32, + *, + loc=None, + ip=None, +) -> Int32: + """Pack four f32 values to FP8 while dividing by one UE8M0 scale.""" + fp32x4 = fp32x4.load() + src_vec4 = fp32x4.ir_value(loc=loc, ip=ip) if hasattr(fp32x4, "ir_value") else fp32x4 + srcs = [ + Float32(vector.extract(src_vec4, [], [index])).ir_value(loc=loc, ip=ip) + for index in range(4) + ] + scale_i32 = Int32(scale_raw_i32).ir_value(loc=loc, ip=ip) + + if cutlass.const_expr(fp8_type is cutlass.Float8E4M3FN): + cvt_instruction = "cvt.rn.satfinite.scaled::n1::ue8m0.e4m3x2.f32" + elif cutlass.const_expr(fp8_type is cutlass.Float8E5M2): + cvt_instruction = "cvt.rn.satfinite.scaled::n1::ue8m0.e5m2x2.f32" + else: + raise ValueError(f"Scaled FP8 conversion does not support {fp8_type}.") + + packed_i32 = llvm.inline_asm( + T.i32(), + [*srcs, scale_i32], + "{\n" + " .reg .b16 lo;\n" + " .reg .b16 hi;\n" + " .reg .b8 s;\n" + " cvt.u8.u32 s, $5;\n" + f" {cvt_instruction} lo, $2, $1, s;\n" + f" {cvt_instruction} hi, $4, $3, s;\n" + " mov.b32 $0, {lo, hi};\n" + "}", + "=r,f,f,f,f,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return Int32(packed_i32) + + +@cute.jit +def cvt_f32x4_to_f8x4_scaled4_pack_i32( + fp32x4: cute.Tensor, + fp8_type, + scales_packed_i32, + *, + loc=None, + ip=None, +) -> Int32: + """Pack four f32 values to FP8 using four independent UE8M0 scales.""" + fp32x4 = fp32x4.load() + src_vec4 = fp32x4.ir_value(loc=loc, ip=ip) if hasattr(fp32x4, "ir_value") else fp32x4 + srcs = [ + Float32(vector.extract(src_vec4, [], [index])).ir_value(loc=loc, ip=ip) + for index in range(4) + ] + scales_i32 = Int32(scales_packed_i32).ir_value(loc=loc, ip=ip) + + if cutlass.const_expr(fp8_type is cutlass.Float8E4M3FN): + cvt_instruction = "cvt.rn.satfinite.scaled::n1::ue8m0.e4m3x2.f32" + elif cutlass.const_expr(fp8_type is cutlass.Float8E5M2): + cvt_instruction = "cvt.rn.satfinite.scaled::n1::ue8m0.e5m2x2.f32" + else: + raise ValueError(f"Scaled FP8 conversion does not support {fp8_type}.") + + packed_i32 = llvm.inline_asm( + T.i32(), + [*srcs, scales_i32], + "{\n" + " .reg .b8 s0, s1, s2, s3;\n" + " .reg .b16 t0, t1, t2, t3;\n" + " .reg .b32 e1, e2, e3, p01, p23;\n" + " shr.u32 e1, $5, 8;\n" + " shr.u32 e2, $5, 16;\n" + " shr.u32 e3, $5, 24;\n" + " cvt.u8.u32 s0, $5;\n" + " cvt.u8.u32 s1, e1;\n" + " cvt.u8.u32 s2, e2;\n" + " cvt.u8.u32 s3, e3;\n" + f" {cvt_instruction} t0, $1, $1, s0;\n" + f" {cvt_instruction} t1, $2, $2, s1;\n" + f" {cvt_instruction} t2, $3, $3, s2;\n" + f" {cvt_instruction} t3, $4, $4, s3;\n" + " mov.b32 p01, {t0, t1};\n" + " mov.b32 p23, {t2, t3};\n" + " prmt.b32 $0, p01, p23, 0x6420;\n" + "}", + "=r,f,f,f,f,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return Int32(packed_i32) + + @dsl_user_op def stg_e8m0_from_f32(addr: Int64, fp32_val: Float32, *, loc=None, ip=None) -> None: """Convert ``fp32_val`` to E8M0 via PTX and store the 1-byte result to global memory. @@ -590,17 +692,9 @@ def stg_e8m0_from_f32(addr: Int64, fp32_val: Float32, *, loc=None, ip=None) -> N @dsl_user_op def stg_e8m0x8_from_f32( addr: Int64, - v0: Float32, - v1: Float32, - v2: Float32, - v3: Float32, - v4: Float32, - v5: Float32, - v6: Float32, - v7: Float32, - *, - loc=None, - ip=None, + v0: Float32, v1: Float32, v2: Float32, v3: Float32, + v4: Float32, v5: Float32, v6: Float32, v7: Float32, + *, loc=None, ip=None, ) -> None: """Convert 8 fp32 values to E8M0 and store them as 8 contiguous bytes in one shot. @@ -613,14 +707,8 @@ def stg_e8m0x8_from_f32( None, [ addr.ir_value(), - v0.ir_value(), - v1.ir_value(), - v2.ir_value(), - v3.ir_value(), - v4.ir_value(), - v5.ir_value(), - v6.ir_value(), - v7.ir_value(), + v0.ir_value(), v1.ir_value(), v2.ir_value(), v3.ir_value(), + v4.ir_value(), v5.ir_value(), v6.ir_value(), v7.ir_value(), ], "{\n" " .reg .b16 p0, p1, p2, p3;\n" @@ -647,13 +735,17 @@ def stg_e8m0x8_from_f32( "cp_reduce_async_bulk_add_bf16_s2g", "cp_reduce_async_bulk_add_u32_s2g", "cvt_f32_to_fp8_to_f32", + "cvt_f32_to_ue8m0_raw_i32", "cvt_f32x4_to_f8x4_pack_i32", + "cvt_f32x4_to_f8x4_scaled4_pack_i32", + "cvt_f32x4_to_f8x4_scaled_pack_i32", "exit", "lds128_v4_b32", "mbarrier_arrive_expect_tx_on_peer", "movmatrix_b16", "nanosleep", "read_clock64", + "read_globaltimer_ns", "red_add_relaxed_sys_f32", "red_add_relaxed_sys_s32", "red_add_relaxed_sys_v2_bf16x2", diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/smem_workspace.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/smem_workspace.py index eb6c51d75cb3..7e5d4af1b466 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/smem_workspace.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/smem_workspace.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Pure-Python SMEM declarations with lifetime-aware overlay placement.""" import dataclasses @@ -12,11 +9,12 @@ from .utils import ( cosize_from_shape_stride_tuples, is_power_of_two, - round_up, row_major_stride, + round_up, validate_static_integer_tuple, ) + SmemRegionKind = Literal["mbarrier", "tensor"] SwizzleSpec = Tuple[int, int, int] @@ -110,7 +108,9 @@ def add_lifetime(self, name: str) -> SmemLifetime: if not name: raise ValueError("An SMEM lifetime needs a non-empty name.") if name in self._lifetime_names: - raise ValueError(f"Duplicate lifetime {name!r} in overlay {self.name!r}.") + raise ValueError( + f"Duplicate lifetime {name!r} in overlay {self.name!r}." + ) lifetime = SmemLifetime(self._workspace, self, name) self._lifetimes.append(lifetime) self._lifetime_names.add(name) @@ -121,7 +121,13 @@ def add_lifetime(self, name: str) -> SmemLifetime: class SmemWorkspace: - """Collect static SMEM declarations and finalize one physical placement.""" + """Collect static SMEM declarations and finalize one physical placement. + + A kernel whose CTAs take different roles builds one workspace per role and lays every resulting offset + table over the same arena. The tables overlap by design: a CTA only reads the table of the role it took, + so a region name a role never registers raises at trace time rather than silently reading another role's + bytes. + """ def __init__( self, @@ -235,17 +241,29 @@ def _make_region( stride = row_major_stride(shape) validate_static_integer_tuple(stride, field_name=f"{name}.stride") if len(shape) != len(stride): - raise ValueError(f"SMEM region {name!r} shape and stride ranks differ.") + raise ValueError( + f"SMEM region {name!r} shape and stride ranks differ." + ) if swizzle is not None: - if len(swizzle) != 3 or not all(isinstance(parameter, int) for parameter in swizzle): - raise TypeError(f"SMEM region {name!r} swizzle must be three Python ints.") + if len(swizzle) != 3 or not all( + isinstance(parameter, int) for parameter in swizzle + ): + raise TypeError( + f"SMEM region {name!r} swizzle must be three Python ints." + ) if swizzle[0] < 0 or swizzle[1] < 0: - raise ValueError(f"SMEM region {name!r} swizzle bits/base must be non-negative.") + raise ValueError( + f"SMEM region {name!r} swizzle bits/base must be non-negative." + ) natural_alignment = max(1, (int(dtype.width) + 7) // 8) - explicit_alignment = natural_alignment if byte_alignment is None else byte_alignment + explicit_alignment = ( + natural_alignment if byte_alignment is None else byte_alignment + ) if not is_power_of_two(explicit_alignment): - raise ValueError(f"SMEM region {name!r} alignment must be a positive power of two.") + raise ValueError( + f"SMEM region {name!r} alignment must be a positive power of two." + ) swizzle_alignment = _swizzle_alignment(swizzle) effective_alignment = max( explicit_alignment, @@ -268,7 +286,10 @@ def _make_region( ) def _claim_region_name(self, region: SmemRegion) -> None: - if region.name in self._region_by_name or region.name in self._overlay_names: + if ( + region.name in self._region_by_name + or region.name in self._overlay_names + ): raise ValueError(f"Duplicate SMEM region {region.name!r}.") self._region_by_name[region.name] = region @@ -287,7 +308,9 @@ def finalize(self, *, max_bytes: Optional[int] = None) -> None: raise RuntimeError("SmemWorkspace.finalize() may only be called once.") offsets, total_bytes = self._build_placement() if max_bytes is not None and total_bytes > max_bytes: - raise ValueError(f"SMEM plan needs {total_bytes} bytes, exceeding {max_bytes} bytes.") + raise ValueError( + f"SMEM plan needs {total_bytes} bytes, exceeding {max_bytes} bytes." + ) self._offset = offsets self._total_bytes = total_bytes self._finalized = True @@ -305,7 +328,9 @@ def _build_placement(self) -> Tuple[Dict[str, int], int]: offsets[declaration.name] = cursor cursor += declaration.nbytes continue - relative_offsets, overlay_alignment, overlay_bytes = self._layout_overlay(declaration) + relative_offsets, overlay_alignment, overlay_bytes = ( + self._layout_overlay(declaration) + ) cursor = round_up(cursor, overlay_alignment) for region_name, relative_offset in relative_offsets.items(): offsets[region_name] = cursor + relative_offset @@ -323,7 +348,9 @@ def _layout_overlay( overlay_bytes = 0 for lifetime in overlay.lifetimes: if not lifetime.regions: - raise ValueError(f"SMEM lifetime {overlay.name}.{lifetime.name} has no regions.") + raise ValueError( + f"SMEM lifetime {overlay.name}.{lifetime.name} has no regions." + ) lifetime_cursor = 0 for region in lifetime.regions: overlay_alignment = max( @@ -369,10 +396,24 @@ class SmemStorage: return SmemStorage + def assert_fits_arena(self, arena_bytes: int) -> None: + """Check this role's placement against the arena every CTA of the launch allocates.""" + self._require_finalized() + if self._total_bytes > arena_bytes: + raise ValueError( + f"SMEM workspace needs {self._total_bytes} bytes but the shared arena is only " + f"{arena_bytes} bytes. The arena is what every CTA allocates, so it cannot be widened for " + "one role without taking the difference away from every other role." + ) + @cute.jit def ptr(self, name: str, smem_base: cute.Pointer) -> cute.Pointer: region = self._region_by_name[name] - swizzle = None if region.swizzle is None else cute.make_swizzle(*region.swizzle) + swizzle = ( + None + if region.swizzle is None + else cute.make_swizzle(*region.swizzle) + ) return cute.make_ptr( region.dtype, smem_base.toint() + self._offset[name], @@ -400,3 +441,4 @@ def _require_finalized(self) -> None: "SmemWorkspace", "SwizzleSpec", ] + diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/software_sync.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/software_sync.py index 135612ecea50..9b1ef4010cb0 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/software_sync.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/software_sync.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Software grid and NVLink synchronization for persistent kernels.""" import cutlass @@ -172,7 +169,10 @@ def arrive_and_wait( signal_base_address = self._signal.toint() signal_byte_offset = Int64(signal_phase * Int32(4)) for rank_round in cutlass.range_constexpr(rank_round_count): - destination_rank = Int32(rank_round * participating_threads) + thread_idx_in_group + destination_rank = ( + Int32(rank_round * participating_threads) + + thread_idx_in_group + ) if destination_rank < Int32(self.world_size): destination_address = self._peer_rank_ptr_mapper.map( signal_base_address, @@ -193,15 +193,12 @@ def arrive_and_wait( scope="gpu", ) local_signal = self._signal + signal_phase - while ( - cute.arch.load( - local_signal, - Int32, - sem="acquire", - scope="sys", - ) - != signal_target - ): + while cute.arch.load( + local_signal, + Int32, + sem="acquire", + scope="sys", + ) != signal_target: pass if cutlass.const_expr(epilogue_grid_sync): diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/topk_reduce.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/topk_reduce.py similarity index 82% rename from tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/topk_reduce.py rename to tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/topk_reduce.py index 5ee767bde5b7..ce4cb6465a0f 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/topk_reduce.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/topk_reduce.py @@ -23,14 +23,16 @@ from typing import ClassVar, Dict, Optional import cuda.bindings.driver as cuda + import cutlass import cutlass.cute as cute -from cutlass._mlir.dialects import llvm from cutlass.cutlass_dsl import Float32, Int32, T +from cutlass._mlir.dialects import llvm + +from ..quant_def import CombineFormat +from .constants import Nvfp4E2M1RcpLimit +from .dsl_helpers import mark_alignment -from .....helpers.constants import Nvfp4E2M1RcpLimit -from .....helpers.dsl_helpers import mark_alignment -from .....quant_def import CombineFormat # --------------------------------------------------------------------------- # fp4 (e2m1) -> fp32 register decode. @@ -198,9 +200,7 @@ def __call__( combine_quant = cute.make_tensor( combine_quant.iterator, - cute.make_layout( - (combine_quant.shape[0], self.num_topk, self.hidden), stride=combine_quant.stride - ), + cute.make_layout((combine_quant.shape[0], self.num_topk, self.hidden), stride=combine_quant.stride), ) reduced_output = cute.make_tensor( reduced_output.iterator, @@ -208,14 +208,11 @@ def __call__( ) if cutlass.const_expr(topk_score is not None): topk_score = cute.make_tensor( - topk_score.iterator, - cute.make_layout((topk_score.shape[0], self.num_topk), stride=topk_score.stride), + topk_score.iterator, cute.make_layout((topk_score.shape[0], self.num_topk), stride=topk_score.stride) ) if cutlass.const_expr(not self.combine_format.is_quantized): - self._reduce_bf16(combine_quant, topk_score, reduced_output).launch( - grid=grid, block=block, stream=stream - ) + self._reduce_bf16(combine_quant, topk_score, reduced_output).launch(grid=grid, block=block, stream=stream) return # The mega kernel hands sf in already as the depth-2 broadcast layout; a @@ -226,8 +223,7 @@ def __call__( sf = cute.make_tensor( combine_sf.iterator, cute.make_layout( - (combine_sf.shape[0], self.num_topk, (sf_vec, self.hidden // sf_vec)), - stride=combine_sf.stride, + (combine_sf.shape[0], self.num_topk, (sf_vec, self.hidden // sf_vec)), stride=combine_sf.stride ), ) else: @@ -251,12 +247,7 @@ def __call__( # -- kernels -------------------------------------------------------------- @cute.kernel - def _reduce_bf16( - self, - combine_output: cute.Tensor, - topk_score: Optional[cute.Tensor], - reduced_output: cute.Tensor, - ): + def _reduce_bf16(self, combine_output: cute.Tensor, topk_score: Optional[cute.Tensor], reduced_output: cute.Tensor): threads = self._threads hidden_per_thread = self.hidden_per_thread hidden_tiles = self.hidden_tiles @@ -269,20 +260,16 @@ def _reduce_bf16( token_idx = worker_idx // hidden_tiles hidden_tile_idx = worker_idx % hidden_tiles - score_dtype = ( - topk_score.dtype if cutlass.const_expr(topk_score is not None) else cutlass.Float32 - ) + score_dtype = topk_score.dtype if cutlass.const_expr(topk_score is not None) else cutlass.Float32 score_reg = cute.make_rmem_tensor((num_topk,), score_dtype) if (not needs_guard) or token_idx < reduced_output.shape[0]: # (token, topk, hidden) -> (topk, hidden_per_thread) - terms = cute.zipped_divide( - combine_output[token_idx, None, None], (num_topk, hidden_per_thread) - )[(None, None), (0, hidden_tile_idx)] - # (token, hidden) -> (hidden_per_thread) - dst = cute.zipped_divide(reduced_output[token_idx, None], (hidden_per_thread,))[ - (None,), (hidden_tile_idx,) + terms = cute.zipped_divide(combine_output[token_idx, None, None], (num_topk, hidden_per_thread))[ + (None, None), (0, hidden_tile_idx) ] + # (token, hidden) -> (hidden_per_thread) + dst = cute.zipped_divide(reduced_output[token_idx, None], (hidden_per_thread,))[(None,), (hidden_tile_idx,)] if cutlass.const_expr(topk_score is not None): if cutlass.const_expr(prefetch): @@ -291,17 +278,13 @@ def _reduce_bf16( for k in cutlass.range_constexpr(num_topk): score_reg[k] = score_dtype(1) - load_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=128 - ) + load_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=128) acc = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float32) for k in cutlass.range_constexpr(0, num_topk, 1): term = cute.make_rmem_tensor((hidden_per_thread,), cutlass.BFloat16) cute.copy( - load_atom, - mark_alignment(terms[k, None], hidden_per_thread * cutlass.BFloat16.width // 8), - term, + load_atom, mark_alignment(terms[k, None], hidden_per_thread * cutlass.BFloat16.width // 8), term ) if cutlass.const_expr(topk_score is not None and not prefetch): score_reg[k] = topk_score[token_idx, Int32(k)] @@ -310,15 +293,12 @@ def _reduce_bf16( for i in cutlass.range_constexpr(0, hidden_per_thread, 2): value_pair = (Float32(term[i]), Float32(term[i + 1])) if cutlass.const_expr(k != 0): + acc[i], acc[i + 1] = cute.arch.fma_packed_f32x2(value_pair, score_pair, (acc[i], acc[i + 1])) + else: + # Include the +0 seed: direct assignment/multiply preserves -0. acc[i], acc[i + 1] = cute.arch.fma_packed_f32x2( - value_pair, score_pair, (acc[i], acc[i + 1]) + value_pair, score_pair, (Float32(0.0), Float32(0.0)) ) - else: - if cutlass.const_expr(topk_score is not None): - acc[i], acc[i + 1] = cute.arch.mul_packed_f32x2(value_pair, score_pair) - else: - acc[i] = value_pair[0] - acc[i + 1] = value_pair[1] out = cute.make_rmem_tensor((hidden_per_thread,), out_dtype) out.store(acc.load().to(out_dtype)) @@ -348,25 +328,21 @@ def _reduce_mxfp8( token_idx = worker_idx // hidden_tiles hidden_tile_idx = worker_idx % hidden_tiles - score_dtype = ( - topk_score.dtype if cutlass.const_expr(topk_score is not None) else cutlass.Float32 - ) + score_dtype = topk_score.dtype if cutlass.const_expr(topk_score is not None) else cutlass.Float32 score_reg = cute.make_rmem_tensor((num_topk,), score_dtype) scale_reg = cute.make_rmem_tensor((num_topk,), cutlass.Float8E8M0FNU) if (not needs_guard) or token_idx < reduced_output.shape[0]: # (token, topk, hidden) -> (topk, hidden_per_thread) - codes = cute.zipped_divide( - combine_quant[token_idx, None, None], (num_topk, hidden_per_thread) - )[(None, None), (0, hidden_tile_idx)] + codes = cute.zipped_divide(combine_quant[token_idx, None, None], (num_topk, hidden_per_thread))[ + (None, None), (0, hidden_tile_idx) + ] # (token, topk, hidden) -> (topk, hidden_per_thread) - sf = cute.zipped_divide( - combine_sf[token_idx, None, None], (num_topk, hidden_per_thread) - )[(None, None), (0, hidden_tile_idx)] - # (token, hidden) -> (hidden_per_thread) - dst = cute.zipped_divide(reduced_output[token_idx, None], (hidden_per_thread,))[ - (None,), (hidden_tile_idx,) + sf = cute.zipped_divide(combine_sf[token_idx, None, None], (num_topk, hidden_per_thread))[ + (None, None), (0, hidden_tile_idx) ] + # (token, hidden) -> (hidden_per_thread) + dst = cute.zipped_divide(reduced_output[token_idx, None], (hidden_per_thread,))[(None,), (hidden_tile_idx,)] if cutlass.const_expr(topk_score is not None): if cutlass.const_expr(prefetch): @@ -375,23 +351,15 @@ def _reduce_mxfp8( for k in cutlass.range_constexpr(num_topk): score_reg[k] = score_dtype(1) if cutlass.const_expr(prefetch): - cute.autovec_copy( - sf[None, 0], scale_reg - ) # one scale per topk slot (stride-0 broadcast) + cute.autovec_copy(sf[None, 0], scale_reg) # one scale per topk slot (stride-0 broadcast) fp8_dtype = self.combine_format.act_dtype - load_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), fp8_dtype, num_bits_per_copy=128 - ) + load_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), fp8_dtype, num_bits_per_copy=128) acc = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float32) for k in cutlass.range_constexpr(0, num_topk, 1): term = cute.make_rmem_tensor((hidden_per_thread,), fp8_dtype) - cute.copy( - load_atom, - mark_alignment(codes[k, None], hidden_per_thread * fp8_dtype.width // 8), - term, - ) + cute.copy(load_atom, mark_alignment(codes[k, None], hidden_per_thread * fp8_dtype.width // 8), term) value = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float32) value.store(term.load().to(cutlass.Float32)) @@ -407,17 +375,12 @@ def _reduce_mxfp8( for i in cutlass.range_constexpr(0, hidden_per_thread, 2): dequant_pair = cute.arch.mul_packed_f32x2((value[i], value[i + 1]), scale_pair) if cutlass.const_expr(k != 0): + acc[i], acc[i + 1] = cute.arch.fma_packed_f32x2(dequant_pair, score_pair, (acc[i], acc[i + 1])) + else: + # Include the +0 seed: direct assignment/multiply preserves -0. acc[i], acc[i + 1] = cute.arch.fma_packed_f32x2( - dequant_pair, score_pair, (acc[i], acc[i + 1]) + dequant_pair, score_pair, (Float32(0.0), Float32(0.0)) ) - else: - if cutlass.const_expr(topk_score is not None): - acc[i], acc[i + 1] = cute.arch.mul_packed_f32x2( - dequant_pair, score_pair - ) - else: - acc[i] = dequant_pair[0] - acc[i + 1] = dequant_pair[1] out = cute.make_rmem_tensor((hidden_per_thread,), out_dtype) out.store(acc.load().to(out_dtype)) @@ -447,25 +410,21 @@ def _reduce_fp4( token_idx = worker_idx // hidden_tiles hidden_tile_idx = worker_idx % hidden_tiles - score_dtype = ( - topk_score.dtype if cutlass.const_expr(topk_score is not None) else cutlass.Float32 - ) + score_dtype = topk_score.dtype if cutlass.const_expr(topk_score is not None) else cutlass.Float32 score_reg = cute.make_rmem_tensor((num_topk,), score_dtype) scale_reg = cute.make_rmem_tensor((num_topk,), cutlass.BFloat16) if (not needs_guard) or token_idx < reduced_output.shape[0]: # (token, topk, hidden) -> (topk, hidden_per_thread) - codes = cute.zipped_divide( - combine_quant[token_idx, None, None], (num_topk, hidden_per_thread) - )[(None, None), (0, hidden_tile_idx)] + codes = cute.zipped_divide(combine_quant[token_idx, None, None], (num_topk, hidden_per_thread))[ + (None, None), (0, hidden_tile_idx) + ] # (token, topk, hidden) -> (topk, hidden_per_thread) - sf = cute.zipped_divide( - combine_sf[token_idx, None, None], (num_topk, hidden_per_thread) - )[(None, None), (0, hidden_tile_idx)] - # (token, hidden) -> (hidden_per_thread) - dst = cute.zipped_divide(reduced_output[token_idx, None], (hidden_per_thread,))[ - (None,), (hidden_tile_idx,) + sf = cute.zipped_divide(combine_sf[token_idx, None, None], (num_topk, hidden_per_thread))[ + (None, None), (0, hidden_tile_idx) ] + # (token, hidden) -> (hidden_per_thread) + dst = cute.zipped_divide(reduced_output[token_idx, None], (hidden_per_thread,))[(None,), (hidden_tile_idx,)] if cutlass.const_expr(topk_score is not None): if cutlass.const_expr(prefetch): @@ -476,19 +435,13 @@ def _reduce_fp4( if cutlass.const_expr(prefetch): cute.autovec_copy(sf[None, 0], scale_reg) - load_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), cutlass.Float4E2M1FN, num_bits_per_copy=64 - ) + load_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Float4E2M1FN, num_bits_per_copy=64) acc = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float32) for k in cutlass.range_constexpr(0, num_topk, 1): term = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float4E2M1FN) cute.copy( - load_atom, - mark_alignment( - codes[k, None], hidden_per_thread * cutlass.Float4E2M1FN.width // 8 - ), - term, + load_atom, mark_alignment(codes[k, None], hidden_per_thread * cutlass.Float4E2M1FN.width // 8), term ) value = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float32) # Dev-only knob (MEGA_F4CVT_USE_MANUAL): manual LUT/PRMT decode vs @@ -513,14 +466,12 @@ def _reduce_fp4( for i in cutlass.range_constexpr(0, hidden_per_thread, 2): dequant_pair = cute.arch.mul_packed_f32x2((value[i], value[i + 1]), scale_pair) if cutlass.const_expr(k != 0): + acc[i], acc[i + 1] = cute.arch.fma_packed_f32x2(dequant_pair, score_pair, (acc[i], acc[i + 1])) + else: + # Include the +0 seed: direct assignment/multiply preserves -0. acc[i], acc[i + 1] = cute.arch.fma_packed_f32x2( - dequant_pair, score_pair, (acc[i], acc[i + 1]) + dequant_pair, score_pair, (Float32(0.0), Float32(0.0)) ) - elif cutlass.const_expr(topk_score is not None): - acc[i], acc[i + 1] = cute.arch.mul_packed_f32x2(dequant_pair, score_pair) - else: - acc[i] = dequant_pair[0] - acc[i + 1] = dequant_pair[1] out = cute.make_rmem_tensor((hidden_per_thread,), out_dtype) out.store(acc.load().to(out_dtype)) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/utils.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/utils.py index 60f768c3f43d..0c3738be2418 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/utils.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/utils.py @@ -1,12 +1,10 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Small integer and layout helpers shared by workspace implementations.""" from typing import Iterable, List, Tuple, Union import cutlass + IntegerType = Union[int, cutlass.Int32, cutlass.Int64, cutlass.Uint32, cutlass.Uint64] @@ -69,9 +67,7 @@ def strides_equal_ignoring_singletons(shape, lhs_stride, rhs_stride) -> bool: return shape == 1 or lhs_stride == rhs_stride -def ordered_stride( - shape: Tuple[int, ...], mem_order: Tuple[int, ...] -) -> Tuple[Tuple[int, ...], int]: +def ordered_stride(shape: Tuple[int, ...], mem_order: Tuple[int, ...]) -> Tuple[Tuple[int, ...], int]: stride = [0] * len(shape) cosize = 1 for mode in sorted(range(len(shape)), key=lambda index: mem_order[index]): diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/__init__.py index 52a7a9daf028..b0d6b3eb2ace 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/__init__.py @@ -1,2 +1,22 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Blackwell inference kernels.""" + +from .mega import ( + BlockScaledSwapAbFc12Extension, + BlockScaledSwapAbFc12Kernel, + BlockScaledSwapAbFc12Mainloop, + BlockScaledSwapAbMegaMoeKernel, + GatedActEpilogueArgs, + SwapABGatedActEpilogue, + TensorRole, +) + + +__all__ = [ + "BlockScaledSwapAbFc12Extension", + "BlockScaledSwapAbFc12Kernel", + "BlockScaledSwapAbFc12Mainloop", + "BlockScaledSwapAbMegaMoeKernel", + "GatedActEpilogueArgs", + "SwapABGatedActEpilogue", + "TensorRole", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/__init__.py index 52a7a9daf028..52aa7758812a 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/__init__.py @@ -1,2 +1,30 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Blackwell MegaMoE inference kernels and issue helpers.""" + +from .block_scaled_swap_ab_fc12_extension import ( + BlockScaledSwapAbFc12Extension, + TensorRole, +) +from .block_scaled_swap_ab_fc12_epilogue import ( + GatedActEpilogueArgs, + SwapABGatedActEpilogue, +) +from .block_scaled_swap_ab_fc12_kernel import ( + BlockScaledSwapAbFc12Kernel, +) +from .block_scaled_swap_ab_fc12_mainloop import ( + BlockScaledSwapAbFc12Mainloop, +) +from .block_scaled_swap_ab_mega_moe_kernel import ( + BlockScaledSwapAbMegaMoeKernel, +) + + +__all__ = [ + "BlockScaledSwapAbFc12Extension", + "BlockScaledSwapAbFc12Kernel", + "BlockScaledSwapAbFc12Mainloop", + "BlockScaledSwapAbMegaMoeKernel", + "GatedActEpilogueArgs", + "SwapABGatedActEpilogue", + "TensorRole", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/custom_mix_cga_helpers.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/custom_mix_cga_helpers.py index 162ec81ab183..32f380cdada0 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/custom_mix_cga_helpers.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/custom_mix_cga_helpers.py @@ -12,16 +12,17 @@ from cutlass._mlir.dialects import cute_nvgpu as _cute_nvgpu_ir from cutlass.cute.nvgpu import cpasync, tcgen05 from cutlass.cute.nvgpu.cpasync.copy import ( - TMA_DESC_PTR_FIELD_NAME, - TMA_MBAR_PTR_FIELD_NAME, - TMA_MCAST_MASK_FIELD_NAME, CopyBulkTensorTileG2SMulticastNonExecTrait, CopyBulkTensorTileG2SMulticastTrait, CopyBulkTensorTileG2SNonExecTrait, CopyBulkTensorTileG2STrait, + TMA_DESC_PTR_FIELD_NAME, + TMA_MBAR_PTR_FIELD_NAME, + TMA_MCAST_MASK_FIELD_NAME, ) from cutlass.cutlass_dsl import Boolean, Int16, Int32 + TmaAtomOrPair = Union[cute.CopyAtom, Tuple[cute.CopyAtom, cute.CopyAtom]] @@ -42,17 +43,12 @@ def make_executable_tma_atom(atom: cute.CopyAtom) -> cute.CopyAtom: def make_executable_tma_atom_or_pair(atom_or_pair: TmaAtomOrPair) -> TmaAtomOrPair: """Materialize one executable atom or a preferred/fallback executable pair.""" if cutlass.const_expr(isinstance(atom_or_pair, tuple)): - return ( - make_executable_tma_atom(atom_or_pair[0]), - make_executable_tma_atom(atom_or_pair[1]), - ) + return (make_executable_tma_atom(atom_or_pair[0]), make_executable_tma_atom(atom_or_pair[1])) return make_executable_tma_atom(atom_or_pair) @cute.jit -def select_tma_atom_for_cluster( - atom_or_pair: TmaAtomOrPair, is_fallback_cluster: Boolean -) -> cute.CopyAtom: +def select_tma_atom_for_cluster(atom_or_pair: TmaAtomOrPair, is_fallback_cluster: Boolean) -> cute.CopyAtom: """Select one preferred/fallback atom after both candidates are materialized.""" if cutlass.const_expr(isinstance(atom_or_pair, tuple)): atom = atom_or_pair[0] @@ -64,34 +60,18 @@ def select_tma_atom_for_cluster( @cute.jit def bind_executable_tma_load_fields( - atom: cute.CopyAtom, - *, - tma_bar_ptr: cute.Pointer, - tma_desc_ptr: Optional[cute.Pointer], - mcast_mask: Optional[Int32], + atom: cute.CopyAtom, *, tma_bar_ptr: cute.Pointer, tma_desc_ptr: Optional[cute.Pointer], mcast_mask: Optional[Int32] ) -> cute.CopyAtom: """Bind per-issue fields to an already executable TMA load atom.""" executable_value = atom._trait.value if cutlass.const_expr(mcast_mask is not None): - mcast_attr = ir.Attribute.parse( - f"#cute_nvgpu.atom_copy_field_tmaload<{TMA_MCAST_MASK_FIELD_NAME}>" - ) - executable_value = _cute_nvgpu_ir.atom_set_value( - executable_value, mcast_attr, Int16(mcast_mask).ir_value() - ) + mcast_attr = ir.Attribute.parse(f"#cute_nvgpu.atom_copy_field_tmaload<{TMA_MCAST_MASK_FIELD_NAME}>") + executable_value = _cute_nvgpu_ir.atom_set_value(executable_value, mcast_attr, Int16(mcast_mask).ir_value()) if cutlass.const_expr(tma_desc_ptr is not None): - descriptor_attr = ir.Attribute.parse( - f"#cute_nvgpu.atom_copy_field_tmaload<{TMA_DESC_PTR_FIELD_NAME}>" - ) - executable_value = _cute_nvgpu_ir.atom_set_value( - executable_value, descriptor_attr, tma_desc_ptr.value - ) - barrier_attr = ir.Attribute.parse( - f"#cute_nvgpu.atom_copy_field_tmaload<{TMA_MBAR_PTR_FIELD_NAME}>" - ) - executable_value = _cute_nvgpu_ir.atom_set_value( - executable_value, barrier_attr, tma_bar_ptr.value - ) + descriptor_attr = ir.Attribute.parse(f"#cute_nvgpu.atom_copy_field_tmaload<{TMA_DESC_PTR_FIELD_NAME}>") + executable_value = _cute_nvgpu_ir.atom_set_value(executable_value, descriptor_attr, tma_desc_ptr.value) + barrier_attr = ir.Attribute.parse(f"#cute_nvgpu.atom_copy_field_tmaload<{TMA_MBAR_PTR_FIELD_NAME}>") + executable_value = _cute_nvgpu_ir.atom_set_value(executable_value, barrier_attr, tma_bar_ptr.value) return cute.CopyAtom(atom.op, atom._trait.__class__(executable_value)) @@ -102,9 +82,7 @@ def _static_vmnk_shape(cta_layout_vmnk: cute.Layout) -> Tuple[int, int, int, int return shape -def _multicast_group_count( - shape_vmnk: Tuple[int, int, int, int], mcast_mode_mn: Tuple[int, int] -) -> int: +def _multicast_group_count(shape_vmnk: Tuple[int, int, int, int], mcast_mode_mn: Tuple[int, int]) -> int: mcast_a_count = shape_vmnk[2] if mcast_mode_mn[0] else 0 mcast_b_count = shape_vmnk[1] if mcast_mode_mn[1] else 0 overlap = 1 if mcast_mode_mn[0] and mcast_mode_mn[1] else 0 @@ -116,10 +94,7 @@ def _multicast_group_count( @cute.jit def _reader_mask( - cta_layout_vmnk: cute.Layout, - cta_coord_vmnk: cute.Coord, - mcast_mode_mn: Tuple[int, int], - mma_cta_count: int, + cta_layout_vmnk: cute.Layout, cta_coord_vmnk: cute.Coord, mcast_mode_mn: Tuple[int, int], mma_cta_count: int ) -> Int32: """Build the general reader union for non-single-sided cluster layouts. @@ -127,23 +102,15 @@ def _reader_mask( """ mask = Int32(0) if cutlass.const_expr(mcast_mode_mn[0]): - mask = mask | Int32( - cpasync.create_tma_multicast_mask(cta_layout_vmnk, cta_coord_vmnk, mcast_mode=2) - ) + mask = mask | Int32(cpasync.create_tma_multicast_mask(cta_layout_vmnk, cta_coord_vmnk, mcast_mode=2)) if cutlass.const_expr(mcast_mode_mn[1]): - mask = mask | Int32( - cpasync.create_tma_multicast_mask(cta_layout_vmnk, cta_coord_vmnk, mcast_mode=1) - ) + mask = mask | Int32(cpasync.create_tma_multicast_mask(cta_layout_vmnk, cta_coord_vmnk, mcast_mode=1)) if cutlass.const_expr(mma_cta_count == 2): peer_coord_vmnk = (cta_coord_vmnk[0] ^ Int32(1), *cta_coord_vmnk[1:]) if cutlass.const_expr(mcast_mode_mn[0]): - mask = mask | Int32( - cpasync.create_tma_multicast_mask(cta_layout_vmnk, peer_coord_vmnk, mcast_mode=2) - ) + mask = mask | Int32(cpasync.create_tma_multicast_mask(cta_layout_vmnk, peer_coord_vmnk, mcast_mode=2)) if cutlass.const_expr(mcast_mode_mn[1]): - mask = mask | Int32( - cpasync.create_tma_multicast_mask(cta_layout_vmnk, peer_coord_vmnk, mcast_mode=1) - ) + mask = mask | Int32(cpasync.create_tma_multicast_mask(cta_layout_vmnk, peer_coord_vmnk, mcast_mode=1)) return mask @@ -156,9 +123,7 @@ def _active_cluster_uses_cluster_sync( preferred_uses_cluster = preferred_cluster_shape_mn[0] * preferred_cluster_shape_mn[1] > 1 fallback_uses_cluster = fallback_cluster_shape_mn[0] * fallback_cluster_shape_mn[1] > 1 cluster_sync_behavior_differs = preferred_uses_cluster != fallback_uses_cluster - return Boolean(preferred_uses_cluster) ^ ( - is_fallback_cluster & Boolean(cluster_sync_behavior_differs) - ) + return Boolean(preferred_uses_cluster) ^ (is_fallback_cluster & Boolean(cluster_sync_behavior_differs)) @cute.jit @@ -233,17 +198,11 @@ def create( or num_mma_consumer_warps <= 0 ): raise ValueError("num_mma_consumer_warps must be a positive Python int.") - if cutlass.const_expr( - len(mcast_mode_mn) != 2 or any(mode not in (0, 1) for mode in mcast_mode_mn) - ): + if cutlass.const_expr(len(mcast_mode_mn) != 2 or any(mode not in (0, 1) for mode in mcast_mode_mn)): raise ValueError("mcast_mode_mn must contain two zero-or-one values.") preferred_shape_vmnk = _static_vmnk_shape(preferred_cta_layout_vmnk) - fallback_layout = ( - preferred_cta_layout_vmnk - if fallback_cta_layout_vmnk is None - else fallback_cta_layout_vmnk - ) + fallback_layout = preferred_cta_layout_vmnk if fallback_cta_layout_vmnk is None else fallback_cta_layout_vmnk fallback_shape_vmnk = _static_vmnk_shape(fallback_layout) if cutlass.const_expr(preferred_shape_vmnk[0] != fallback_shape_vmnk[0]): raise ValueError("Preferred and fallback layouts must use the same MMA V extent.") @@ -256,21 +215,15 @@ def create( cta_group = tcgen05.CtaGroup.TWO if mma_cta_count == 2 else tcgen05.CtaGroup.ONE is_true_mixed = preferred_shape_vmnk != fallback_shape_vmnk if cutlass.const_expr(is_true_mixed and is_fallback_cluster is None): - raise ValueError( - "is_fallback_cluster is required when preferred and fallback layouts differ." - ) + raise ValueError("is_fallback_cluster is required when preferred and fallback layouts differ.") preferred_group_count = _multicast_group_count(preferred_shape_vmnk, mcast_mode_mn) fallback_group_count = _multicast_group_count(fallback_shape_vmnk, mcast_mode_mn) preferred_arrival_count = preferred_group_count * num_mma_consumer_warps fallback_arrival_count = fallback_group_count * num_mma_consumer_warps - preferred_cluster_size = ( - preferred_shape_vmnk[0] * preferred_shape_vmnk[1] * preferred_shape_vmnk[2] - ) - fallback_cluster_size = ( - fallback_shape_vmnk[0] * fallback_shape_vmnk[1] * fallback_shape_vmnk[2] - ) + preferred_cluster_size = preferred_shape_vmnk[0] * preferred_shape_vmnk[1] * preferred_shape_vmnk[2] + fallback_cluster_size = fallback_shape_vmnk[0] * fallback_shape_vmnk[1] * fallback_shape_vmnk[2] tracks_both_operand_families = mcast_mode_mn == (1, 1) is_single_sided = tracks_both_operand_families and ( (preferred_shape_vmnk[1] == fallback_shape_vmnk[1] == 1) @@ -285,9 +238,7 @@ def create( if cutlass.const_expr(is_single_sided): consumer_mask = Int32((1 << fallback_cluster_size) - 1) else: - consumer_mask = _reader_mask( - fallback_layout, cta_coord_vmnk, mcast_mode_mn, mma_cta_count - ) + consumer_mask = _reader_mask(fallback_layout, cta_coord_vmnk, mcast_mode_mn, mma_cta_count) else: active_arrival_count = Int32(preferred_arrival_count) if cutlass.const_expr(is_single_sided): @@ -297,9 +248,7 @@ def create( preferred_cta_layout_vmnk, cta_coord_vmnk, mcast_mode_mn, mma_cta_count ) elif cutlass.const_expr(not is_single_sided): - consumer_mask = _reader_mask( - preferred_cta_layout_vmnk, cta_coord_vmnk, mcast_mode_mn, mma_cta_count - ) + consumer_mask = _reader_mask(preferred_cta_layout_vmnk, cta_coord_vmnk, mcast_mode_mn, mma_cta_count) full_sync_object = pipeline.MbarrierArray( barrier_storage=barrier_storage.align(min_align=8), diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/__init__.py index 52a7a9daf028..b84f4d01b814 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/__init__.py @@ -1,2 +1 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Rubin kernel implementations.""" diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/__init__.py index 52a7a9daf028..9f3a0e63c832 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/__init__.py @@ -1,2 +1,30 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Rubin inference kernels.""" + +from .mega import ( + BlockScaledSwapAbFc12Extension, + BlockScaledSwapAbFc12Mainloop, + BlockScaledSwapAbFc12MainloopGenSpecialized, + BlockScaledSwapAbGenphaseMoeKernel, + BlockScaledSwapAbMegaMoeKernel, + GatedActEpilogueArgs, + GenphaseTokenComm, + GenphaseTokenCommArgs, + SwapABGatedActEpilogue, + TensorRole, + dynamic_mainloop, +) + + +__all__ = [ + "BlockScaledSwapAbFc12Extension", + "BlockScaledSwapAbFc12Mainloop", + "BlockScaledSwapAbFc12MainloopGenSpecialized", + "BlockScaledSwapAbGenphaseMoeKernel", + "BlockScaledSwapAbMegaMoeKernel", + "GenphaseTokenComm", + "GenphaseTokenCommArgs", + "GatedActEpilogueArgs", + "SwapABGatedActEpilogue", + "TensorRole", + "dynamic_mainloop", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py index 33b754f7b9cd..8057166cabff 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py @@ -10,29 +10,26 @@ """ import dataclasses +import math from typing import Any, Callable, ClassVar, List, Literal, Optional, Tuple, Type, Union import cutlass import cutlass.cute as cute -import cutlass.pipeline as pipeline +from cutlass.cutlass_dsl import Int64, T +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.typing import AddressSpace import cutlass.utils as utils +import cutlass.pipeline as pipeline import cutlass.utils.blackwell_helpers as sm100_utils + from cutlass._mlir import ir from cutlass._mlir.dialects import llvm, vector -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass.cute.typing import AddressSpace -from cutlass.cutlass_dsl import Int64, T -from .....api import ( - ImplDesc, - KernelComponent, - OptionalRequirement, - ProblemDesc, - StaticOrRuntimeIntegerType, -) from .....communication.nvlink_domain.symmetric_buffer import SymmetricBufferDevice from .....communication.token_protocol import TokenSrcMetadata -from .....helpers.constants import Fp8E4M3RcpLimit, Fp8E5M2RcpLimit, Fp32Max, Nvfp4E2M1RcpLimit +from .....quant_def import CombineFormat, QuantKind +from .....api import ImplDesc, KernelComponent, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType +from .....helpers.constants import Nvfp4E2M1RcpLimit, Fp8E4M3RcpLimit, Fp8E5M2RcpLimit, Fp32Max from .....helpers.cute_py_helpers import tcgen05_block_scaled_acc_dtype from .....helpers.dsl_helpers import mark_alignment from .....helpers.flag_batch import GpuReleaseFlagBatchTracker @@ -44,10 +41,8 @@ red_add_relaxed_sys_v2_bf16x2, ) from .....helpers.smem_workspace import SmemRegion, SmemWorkspace -from .....quant_def import CombineFormat, QuantKind from ....function_mapping import CoordinateSpace, FunctionMapping -from ....schedulers.base import SchedulerConsumer -from ....schedulers.fc12_mapping import BlockPhase, SwapAbFc12WorkTileInfo +from ....schedulers import BlockPhase, SchedulerConsumer, SwapAbFc12WorkTileInfo from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension @@ -82,43 +77,27 @@ class QuantImpl: warp_idx: Optional[Any] = None pair_exchange_barrier: Optional[Any] = None - _directions: ClassVar[Tuple[str, ...]] = ( - "regs_in_thread", - "threads_with_the_same_reg", - "regs_in_pair_threads", - ) + _directions: ClassVar[Tuple[str, ...]] = ("regs_in_thread", "threads_with_the_same_reg", "regs_in_pair_threads") # -- config / validation -------------------------------------------------- def __post_init__(self): if isinstance(self.quant_kind, CombineFormat): if not self.quant_kind.is_quantized: - raise ValueError( - f"QuantImpl combine path needs a quantized CombineFormat, got {self.quant_kind}." - ) + raise ValueError(f"QuantImpl combine path needs a quantized CombineFormat, got {self.quant_kind}.") elif not isinstance(self.quant_kind, QuantKind): - raise ValueError( - f"quant_kind must be a QuantKind or a CombineFormat, got {self.quant_kind!r}." - ) + raise ValueError(f"quant_kind must be a QuantKind or a CombineFormat, got {self.quant_kind!r}.") if self.sf_vec_direction not in self._directions: - raise ValueError( - f"sf_vec_direction must be one of {self._directions}, got {self.sf_vec_direction!r}." - ) + raise ValueError(f"sf_vec_direction must be one of {self._directions}, got {self.sf_vec_direction!r}.") if self.sf_vec_direction == "threads_with_the_same_reg" and self.lane_idx is None: raise ValueError("across-lane quant needs lane_idx for the CREDUX half-warp predicate.") if self.sf_vec_direction == "regs_in_pair_threads": if isinstance(self.quant_kind, CombineFormat): - raise ValueError( - "The paired direction is orthodox-only; combine blocks never straddle warps." - ) + raise ValueError("The paired direction is orthodox-only; combine blocks never straddle warps.") if self.sf_vec_size % 2 != 0: - raise ValueError( - f"The paired direction needs an even sf_vec, got {self.sf_vec_size}." - ) + raise ValueError(f"The paired direction needs an even sf_vec, got {self.sf_vec_size}.") if self.lane_idx is None or self.warp_idx is None or self.pair_exchange_barrier is None: - raise ValueError( - "The paired direction needs lane_idx, warp_idx and pair_exchange_barrier." - ) + raise ValueError("The paired direction needs lane_idx, warp_idx and pair_exchange_barrier.") elif not isinstance(self.quant_kind, CombineFormat) and self.sf_vec_size != 16: # A wider orthodox block cannot be reduced inside one thread: see the paired direction. raise ValueError(f"Orthodox sf_vec {self.sf_vec_size} needs the paired direction.") @@ -159,20 +138,14 @@ def _data_rcp_limit(self) -> float: @property def _block_share_per_thread(self) -> int: """How many of a block's elements this thread holds.""" - return ( - self.sf_vec_size // 2 - if self.sf_vec_direction == "regs_in_pair_threads" - else self.sf_vec_size - ) + return self.sf_vec_size // 2 if self.sf_vec_direction == "regs_in_pair_threads" else self.sf_vec_size # -- dispatch ------------------------------------------------------------- @cute.jit def __call__(self, prequant_reg: cute.Tensor, *, norm_const=None, smem_intermediate=None): if cutlass.const_expr(cute.size(prequant_reg) % self._block_share_per_thread != 0): - raise ValueError( - "prequant_reg size must be divisible by this thread's share of a block." - ) + raise ValueError("prequant_reg size must be divisible by this thread's share of a block.") # Combine quantizes fc2's bf16 reorder regs; orthodox the fp32 swiglu. expected_in = cutlass.BFloat16 if self._is_combine else cutlass.Float32 if cutlass.const_expr(prequant_reg.element_type is not expected_in): @@ -181,9 +154,7 @@ def __call__(self, prequant_reg: cute.Tensor, *, norm_const=None, smem_intermedi ) needs_smem = cutlass.const_expr(self.sf_vec_direction == "regs_in_pair_threads") if cutlass.const_expr(needs_smem != (smem_intermediate is not None)): - raise ValueError( - "smem_intermediate must be supplied for the paired direction and only for it." - ) + raise ValueError("smem_intermediate must be supplied for the paired direction and only for it.") if cutlass.const_expr(not self._is_combine): if cutlass.const_expr(self.sf_vec_direction == "regs_in_pair_threads"): return self.mx_quant_regs_in_pair_threads_impl(prequant_reg, smem_intermediate) @@ -255,14 +226,10 @@ def mx_quant_regs_in_pair_threads_impl( if cutlass.const_expr(smem_intermediate.stride[0] != 1): # Thread-major would still be correct but would silently split the exchange into # per-block scalar accesses. - raise ValueError( - "The paired amax exchange must be block-major to stay one vector access." - ) + raise ValueError("The paired amax exchange must be block-major to stay one vector access.") exchange_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), - cutlass.Float32, - num_bits_per_copy=n_blocks * cutlass.Float32.width, + cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=n_blocks * cutlass.Float32.width ) # Which warps pair up is epilogue knowledge: warp w owns one contiguous run of # intermediate outputs, so a block that is twice that run wide joins adjacent warps. @@ -273,20 +240,12 @@ def mx_quant_regs_in_pair_threads_impl( half_amax = cute.make_rmem_tensor((n_blocks,), cutlass.Float32) for vec_block_idx in cutlass.range_constexpr(n_blocks): half_amax[vec_block_idx] = self._amax_thread_fp32(in_blocks[None, vec_block_idx]) - cute.copy( - exchange_atom, - cute.coalesce(half_amax), - cute.coalesce(smem_intermediate[None, own_thread]), - ) + cute.copy(exchange_atom, cute.coalesce(half_amax), cute.coalesce(smem_intermediate[None, own_thread])) self.pair_exchange_barrier.arrive_and_wait() partner_amax = cute.make_rmem_tensor((n_blocks,), cutlass.Float32) - cute.copy( - exchange_atom, - cute.coalesce(smem_intermediate[None, partner_thread]), - cute.coalesce(partner_amax), - ) + cute.copy(exchange_atom, cute.coalesce(smem_intermediate[None, partner_thread]), cute.coalesce(partner_amax)) data_blocks = [] sf_values = [] @@ -294,17 +253,13 @@ def mx_quant_regs_in_pair_threads_impl( block_amax = cute.arch.fmax(half_amax[vec_block_idx], partner_amax[vec_block_idx]) scale_e8m0, scale_f32 = self._e8m0(block_amax) sf_values.append(scale_e8m0) - data_blocks.append( - self._scale_to_data_ssa(in_blocks[None, vec_block_idx], self._enc_mxfp8(scale_f32)) - ) + data_blocks.append(self._scale_to_data_ssa(in_blocks[None, vec_block_idx], self._enc_mxfp8(scale_f32))) self._store_packed_blocks(data, data_blocks) sf.store(self._values_to_ssa(sf_values, self.scale_dtype)) return data, sf @cute.jit - def nvfp4_combine_quant_regs_in_thread_impl( - self, prequant_reg: cute.Tensor - ) -> Tuple[cute.Tensor, cute.Tensor]: + def nvfp4_combine_quant_regs_in_thread_impl(self, prequant_reg: cute.Tensor) -> Tuple[cute.Tensor, cute.Tensor]: # bf16 in -> e2m1 data + per-16 bf16 amax. amax found on bf16 (packed). vec = self.sf_vec_size n_blocks = cute.size(prequant_reg) // vec @@ -361,9 +316,7 @@ def nvfp4_combine_quant_threads_with_the_same_reg_impl( return data, selected_sf @cute.jit - def mxfp8_combine_quant_regs_in_thread_impl( - self, prequant_reg: cute.Tensor - ) -> Tuple[cute.Tensor, cute.Tensor]: + def mxfp8_combine_quant_regs_in_thread_impl(self, prequant_reg: cute.Tensor) -> Tuple[cute.Tensor, cute.Tensor]: # bf16 in -> e4m3/e5m2 data + per-32 e8m0. amax found on bf16 (packed). vec = self.sf_vec_size n_blocks = cute.size(prequant_reg) // vec @@ -439,8 +392,7 @@ def _concat_blocks_ssa(self, blocks, dtype: Type[cutlass.Numeric]) -> cute.Tenso @cute.jit def _values_to_ssa(self, values, dtype: Type[cutlass.Numeric]) -> cute.TensorSSA: vec = vector.from_elements( - T.vector(len(values), dtype.mlir_type), - [values[i].ir_value() for i in range(len(values))], + T.vector(len(values), dtype.mlir_type), [values[i].ir_value() for i in range(len(values))] ) return cute.TensorSSA(vec, (len(values),), dtype) @@ -484,9 +436,7 @@ def max_abs(lhs: cutlass.Float32, rhs: cutlass.Float32) -> cutlass.Float32: acc = block[0] for elem_idx in cutlass.range_constexpr(1, cute.size(block)): acc = max_abs(acc, block[elem_idx]) - mag_bits = cutlass.Int32( - llvm.bitcast(T.i32(), cutlass.Float32(acc).ir_value()) - ) & cutlass.Int32(0x7FFFFFFF) + mag_bits = cutlass.Int32(llvm.bitcast(T.i32(), cutlass.Float32(acc).ir_value())) & cutlass.Int32(0x7FFFFFFF) return cutlass.Float32(llvm.bitcast(T.f32(), mag_bits.ir_value())) @cute.jit @@ -592,15 +542,12 @@ class _TmemTranspose16x32Core: @staticmethod def _tmem_layout(num_lanes: int, num_cols: int) -> cute.Layout: return cute.make_layout( - (((num_lanes, num_cols), 1),), - stride=(((_TmemTranspose16x32Core._TmemRowStride, 1), 0),), + (((num_lanes, num_cols), 1),), stride=(((_TmemTranspose16x32Core._TmemRowStride, 1), 0),) ) @staticmethod def _rmem_copy_view(rmem: cute.Tensor, num_regs: int, offset: int = 0) -> cute.Tensor: - return cute.make_tensor( - rmem.iterator + offset, cute.make_layout((((num_regs,), 1),), stride=(((1,), 0),)) - ) + return cute.make_tensor(rmem.iterator + offset, cute.make_layout((((num_regs,), 1),), stride=(((1,), 0),))) @staticmethod def load_subtile_raw_acc( @@ -646,31 +593,19 @@ def load_subtile_raw_acc( # offsets here would wrap them as SSA values that cute treats as # alignment-unknown, tripping the atom's verifier. first_top_view = cute.make_tensor(ptr, _TmemTranspose16x32Core._tmem_layout(16, 32)) - first_bot_view = cute.make_tensor( - ptr + half_lane_off, _TmemTranspose16x32Core._tmem_layout(16, 32) - ) + first_bot_view = cute.make_tensor(ptr + half_lane_off, _TmemTranspose16x32Core._tmem_layout(16, 32)) second_top_view = cute.make_tensor(ptr + 32, _TmemTranspose16x32Core._tmem_layout(16, 32)) - second_bot_view = cute.make_tensor( - ptr + 32 + half_lane_off, _TmemTranspose16x32Core._tmem_layout(16, 32) - ) + second_bot_view = cute.make_tensor(ptr + 32 + half_lane_off, _TmemTranspose16x32Core._tmem_layout(16, 32)) first_top = cute.make_rmem_tensor((16,), _TmemTranspose16x32Core._io_dtype) first_bot = cute.make_rmem_tensor((16,), _TmemTranspose16x32Core._io_dtype) second_top = cute.make_rmem_tensor((16,), _TmemTranspose16x32Core._io_dtype) second_bot = cute.make_rmem_tensor((16,), _TmemTranspose16x32Core._io_dtype) - cute.copy( - atom_ld16x64, first_top_view, _TmemTranspose16x32Core._rmem_copy_view(first_top, 16) - ) - cute.copy( - atom_ld16x64, first_bot_view, _TmemTranspose16x32Core._rmem_copy_view(first_bot, 16) - ) - cute.copy( - atom_ld16x64, second_top_view, _TmemTranspose16x32Core._rmem_copy_view(second_top, 16) - ) - cute.copy( - atom_ld16x64, second_bot_view, _TmemTranspose16x32Core._rmem_copy_view(second_bot, 16) - ) + cute.copy(atom_ld16x64, first_top_view, _TmemTranspose16x32Core._rmem_copy_view(first_top, 16)) + cute.copy(atom_ld16x64, first_bot_view, _TmemTranspose16x32Core._rmem_copy_view(first_bot, 16)) + cute.copy(atom_ld16x64, second_top_view, _TmemTranspose16x32Core._rmem_copy_view(second_top, 16)) + cute.copy(atom_ld16x64, second_bot_view, _TmemTranspose16x32Core._rmem_copy_view(second_bot, 16)) return (first_top, first_bot, second_top, second_bot) @@ -703,21 +638,11 @@ def __init__(self, tmem_ptr, region: int, reg_tensor: Optional[cute.Tensor] = No self._tmem_dst_top = cute.make_tensor(dst_ptr, self._tmem_layout(16, 16)) self._tmem_dst_bot = cute.make_tensor(dst_ptr + half_lane_off, self._tmem_layout(16, 16)) - self._atom_ld16x64 = cute.make_copy_atom( - tcgen05.Ld16x64bOp(tcgen05.Repetition.x16), self._io_dtype - ) - self._atom_st16x128 = cute.make_copy_atom( - tcgen05.St16x128bOp(tcgen05.Repetition.x8), self._io_dtype - ) - self._atom_st32x32 = cute.make_copy_atom( - tcgen05.St32x32bOp(tcgen05.Repetition.x16), self._io_dtype - ) - self._atom_ld16x256 = cute.make_copy_atom( - tcgen05.Ld16x256bOp(tcgen05.Repetition.x2), self._io_dtype - ) - self._atom_ld16x128 = cute.make_copy_atom( - tcgen05.Ld16x128bOp(tcgen05.Repetition.x4), self._io_dtype - ) + self._atom_ld16x64 = cute.make_copy_atom(tcgen05.Ld16x64bOp(tcgen05.Repetition.x16), self._io_dtype) + self._atom_st16x128 = cute.make_copy_atom(tcgen05.St16x128bOp(tcgen05.Repetition.x8), self._io_dtype) + self._atom_st32x32 = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x16), self._io_dtype) + self._atom_ld16x256 = cute.make_copy_atom(tcgen05.Ld16x256bOp(tcgen05.Repetition.x2), self._io_dtype) + self._atom_ld16x128 = cute.make_copy_atom(tcgen05.Ld16x128bOp(tcgen05.Repetition.x4), self._io_dtype) self._src_regs = cute.make_rmem_tensor((16,), self._io_dtype) # ``output`` is a bare (16,) RMEM fragment; its (lane_idx, elem_idx) @@ -772,18 +697,10 @@ def r2_store(self) -> None: # -- R3 ------------------------------------------------------------------ def r3_load_top(self) -> None: - cute.copy( - self._atom_ld16x256, - self._tmem_dst_top, - self._rmem_copy_view(self._src_regs, 8, offset=0), - ) + cute.copy(self._atom_ld16x256, self._tmem_dst_top, self._rmem_copy_view(self._src_regs, 8, offset=0)) def r3_load_bot(self) -> None: - cute.copy( - self._atom_ld16x256, - self._tmem_dst_bot, - self._rmem_copy_view(self._src_regs, 8, offset=8), - ) + cute.copy(self._atom_ld16x256, self._tmem_dst_bot, self._rmem_copy_view(self._src_regs, 8, offset=8)) def r3_perm(self) -> None: for r in range(16): @@ -795,18 +712,10 @@ def r3_store(self) -> None: # -- R4 ------------------------------------------------------------------ def r4_load_top(self) -> None: - cute.copy( - self._atom_ld16x128, - self._tmem_dst_top, - self._rmem_copy_view(self._src_regs, 8, offset=0), - ) + cute.copy(self._atom_ld16x128, self._tmem_dst_top, self._rmem_copy_view(self._src_regs, 8, offset=0)) def r4_load_bot(self) -> None: - cute.copy( - self._atom_ld16x128, - self._tmem_dst_bot, - self._rmem_copy_view(self._src_regs, 8, offset=8), - ) + cute.copy(self._atom_ld16x128, self._tmem_dst_bot, self._rmem_copy_view(self._src_regs, 8, offset=8)) def r4_perm(self) -> None: for r in range(16): @@ -866,10 +775,7 @@ class TmemTranspose32x32Inplace: """ def __init__( - self, - tmem_ptr, - reg_tensor_top: Optional[cute.Tensor] = None, - reg_tensor_bot: Optional[cute.Tensor] = None, + self, tmem_ptr, reg_tensor_top: Optional[cute.Tensor] = None, reg_tensor_bot: Optional[cute.Tensor] = None ) -> None: if (reg_tensor_top is None) != (reg_tensor_bot is None): raise ValueError( @@ -962,26 +868,19 @@ class TmemTranspose32x64B16Movm: @staticmethod def _tmem_layout(num_lanes: int, num_cols: int) -> cute.Layout: return cute.make_layout( - (((num_lanes, num_cols), 1),), - stride=(((TmemTranspose32x64B16Movm._tmem_row_stride, 1), 0),), + (((num_lanes, num_cols), 1),), stride=(((TmemTranspose32x64B16Movm._tmem_row_stride, 1), 0),) ) @staticmethod def _rmem_copy_view(rmem: cute.Tensor, num_regs: int, offset: int = 0) -> cute.Tensor: - return cute.make_tensor( - rmem.iterator + offset, cute.make_layout((((num_regs,), 1),), stride=(((1,), 0),)) - ) + return cute.make_tensor(rmem.iterator + offset, cute.make_layout((((num_regs,), 1),), stride=(((1,), 0),))) @cute.jit def __init__(self, tmem_ptr, reg_tensor: cute.Tensor) -> None: if cutlass.const_expr(reg_tensor.element_type is not cutlass.BFloat16): - raise TypeError( - f"{type(self).__name__} expects BF16 input after f2fp, got {reg_tensor.element_type}." - ) + raise TypeError(f"{type(self).__name__} expects BF16 input after f2fp, got {reg_tensor.element_type}.") if cutlass.const_expr(cute.size(reg_tensor) != 64): - raise ValueError( - f"{type(self).__name__} expects 64 BF16 elements, got {cute.size(reg_tensor)}." - ) + raise ValueError(f"{type(self).__name__} expects 64 BF16 elements, got {cute.size(reg_tensor)}.") movm_words = movmatrix_b16(cute.recast_tensor(reg_tensor, cutlass.Int32)) self._store_words = cute.make_rmem_tensor(movm_words.layout, movm_words.element_type) @@ -992,20 +891,14 @@ def __init__(self, tmem_ptr, reg_tensor: cute.Tensor) -> None: self._tmem_top = cute.make_tensor(tmem_ptr, self._tmem_layout(16, 32)) self._tmem_bottom = cute.make_tensor(tmem_ptr + half_lane_offset, self._tmem_layout(16, 32)) self._tmem_full = cute.make_tensor(tmem_ptr, self._tmem_layout(32, 32)) - self._store_atom = cute.make_copy_atom( - tcgen05.St16x128bOp(tcgen05.Repetition.x8), cutlass.Float32 - ) - self._load_atom = cute.make_copy_atom( - tcgen05.Ld32x32bOp(tcgen05.Repetition.x32), cutlass.Float32 - ) + self._store_atom = cute.make_copy_atom(tcgen05.St16x128bOp(tcgen05.Repetition.x8), cutlass.Float32) + self._load_atom = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x32), cutlass.Float32) @cute.jit def __call__(self) -> cute.Tensor: movm_words_f32 = cute.recast_tensor(self._store_words, cutlass.Float32) cute.copy(self._store_atom, self._rmem_copy_view(movm_words_f32, 16), self._tmem_top) - cute.copy( - self._store_atom, self._rmem_copy_view(movm_words_f32, 16, offset=16), self._tmem_bottom - ) + cute.copy(self._store_atom, self._rmem_copy_view(movm_words_f32, 16, offset=16), self._tmem_bottom) output_words = cute.make_rmem_tensor((32,), cutlass.Float32) cute.copy(self._load_atom, self._tmem_full, self._rmem_copy_view(output_words, 32)) @@ -1020,8 +913,7 @@ class GatedActEpilogueArgs: fc2_alpha: Optional[cute.Tensor] fc1_norm_const: Optional[cute.Tensor] # ----------------------------------- - # MoE domain (token, topk), deepgemm graph only? for transformer graph, we want reduce kernel to perform the - # score mul. + # MoE domain (token, topk), deepgemm graph only? for transformer graph, we want reduce kernel to perform the score mul. topk_scores: Optional[cute.Tensor] @@ -1102,19 +994,13 @@ def _resolve_situ_betas(problem_desc: ProblemDesc) -> Tuple[Optional[float], Opt beta = problem_desc.get("situ_beta") linear_beta = problem_desc.get("situ_linear_beta") if (beta is None) != (linear_beta is None): - raise ValueError( - f"situ_beta and situ_linear_beta must be set together, got {beta} and {linear_beta}." - ) + raise ValueError(f"situ_beta and situ_linear_beta must be set together, got {beta} and {linear_beta}.") if beta is None: return None, None if beta <= 0 or linear_beta <= 0: - raise ValueError( - f"SiTU beta parameters must be positive, got {beta} and {linear_beta}." - ) + raise ValueError(f"SiTU beta parameters must be positive, got {beta} and {linear_beta}.") if problem_desc["gate_up_clamp"] is not None: - raise ValueError( - "SiTU does not support gate_up_clamp; the two activation variants are exclusive." - ) + raise ValueError("SiTU does not support gate_up_clamp; the two activation variants are exclusive.") if problem_desc.get("swiglu_alpha") is not None: raise ValueError("SiTU does not support SwiGLUBias constants.") return float(beta), float(linear_beta) @@ -1147,17 +1033,11 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: if self.communication_enabled: for field_name in ("reduce_topk_in_kernel", "token_back_push_data"): if field_name not in impl_desc: - raise KeyError( - f"Communication-enabled Epilogue requires ImplDesc field {field_name!r}." - ) + raise KeyError(f"Communication-enabled Epilogue requires ImplDesc field {field_name!r}.") self.reduce_topk_in_kernel = impl_desc.get("reduce_topk_in_kernel", False) self.token_back_push_data = impl_desc.get("token_back_push_data", False) - if not self.communication_enabled and ( - self.reduce_topk_in_kernel or self.token_back_push_data - ): - raise ValueError( - "A communication-disabled Epilogue cannot enable communication policies." - ) + if not self.communication_enabled and (self.reduce_topk_in_kernel or self.token_back_push_data): + raise ValueError("A communication-disabled Epilogue cannot enable communication policies.") # FC1 emits what FC2 consumes as its activation, in the kind's own scale format. self.fc1_output_dtype = self.quant_kind.activation_dtype @@ -1177,9 +1057,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.reduce_topk_in_epilogue = self.reduce_topk_in_kernel and not self.token_back_push_data if not 1 <= self.fc1_epi_flag_batch <= 32 or not 1 <= self.fc2_epi_flag_batch <= 32: raise ValueError("Epilogue flag batch sizes must be in [1, 32].") - self.cluster_tile_intermediate_downproj = ( - self._EpilogueFc1IntermediateDownTileSize * self.cluster_shape_mn[0] - ) + self.cluster_tile_intermediate_downproj = self._EpilogueFc1IntermediateDownTileSize * self.cluster_shape_mn[0] atom_thr_size = 2 if self.use_2cta_instrs else 1 self.cta_tile_m = self._EpilogueFc2HiddenTileSize @@ -1199,17 +1077,10 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: ) self.num_accumulator_tmem_cols = tmem_plan.accumulator_columns self.overlapping_accum = self.accumulator_overlap_columns > 0 - self.accumulator_shape = ( - self.cta_tile_m, - self.cta_tile_n, - tmem_plan.accumulator_stage_count, - ) + self.accumulator_shape = (self.cta_tile_m, self.cta_tile_n, tmem_plan.accumulator_stage_count) self.accumulator_stride = (1 << 16, 1, tmem_plan.accumulator_stage_stride_columns) - if ( - isinstance(self.hidden_size, int) - and self.hidden_size % (self.cta_tile_m * self.cluster_shape_mn[0]) == 0 - ): + if isinstance(self.hidden_size, int) and self.hidden_size % (self.cta_tile_m * self.cluster_shape_mn[0]) == 0: self.fc2_hidden_needs_predicate: bool = False else: self.fc2_hidden_needs_predicate: bool = True @@ -1223,10 +1094,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: # One staging stage per token subtile, each an (epi_tile_n, epi_tile_m // 2) quantized tile. self.fc1_staging_stage_bytes = ( - self._EpilogueTokenTileSize - * self._EpilogueFc1IntermediateDownTileSize - * self.fc1_output_dtype.width - // 8 + self._EpilogueTokenTileSize * self._EpilogueFc1IntermediateDownTileSize * self.fc1_output_dtype.width // 8 ) self.fc1_staging_bytes = self.subtile_cnt * self.fc1_staging_stage_bytes # The amax exchange plane is (token chunk, epilogue thread), block-major so each thread's @@ -1235,32 +1103,21 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: # bytes -- see the lifetime argument in fc1_quant. self.fc1_amax_token_chunks = self._EpilogueTokenTileSize // 32 self.fc1_amax_slot_count = ( - self.fc1_amax_token_chunks * self._EpilogueWarpCnt * 32 - if self.needs_pair_amax_exchange - else 0 + self.fc1_amax_token_chunks * self._EpilogueWarpCnt * 32 if self.needs_pair_amax_exchange else 0 ) if self.fc1_amax_slot_count * 4 > self.fc1_staging_stage_bytes: raise ValueError("The paired amax exchange does not fit in one fc1 staging stage.") requested_fc2_tma_stages = impl_desc.get("fc2_tma_stages") - if ( - requested_fc2_tma_stages is not None - and not 1 <= requested_fc2_tma_stages <= self.subtile_cnt - ): - raise ValueError( - f"fc2_tma_stages must be in [1, {self.subtile_cnt}], got {requested_fc2_tma_stages}." - ) + if requested_fc2_tma_stages is not None and not 1 <= requested_fc2_tma_stages <= self.subtile_cnt: + raise ValueError(f"fc2_tma_stages must be in [1, {self.subtile_cnt}], got {requested_fc2_tma_stages}.") if self.fc2_use_bulk: single_stage_region = self._make_fc2_single_stage_region() if self.fc2_use_ublk and single_stage_region.nbytes % 16 != 0: raise ValueError("Each FC2 UBLK staging stage must occupy a multiple of 16 bytes.") # Additional stages trade mainloop SMEM for store overlap. - self.fc2_tma_stages = ( - requested_fc2_tma_stages if requested_fc2_tma_stages is not None else 1 - ) - self.fc2_staging_spec: Optional[SmemRegion] = self._make_fc2_staging_region( - self.fc2_tma_stages - ) + self.fc2_tma_stages = requested_fc2_tma_stages if requested_fc2_tma_stages is not None else 1 + self.fc2_staging_spec: Optional[SmemRegion] = self._make_fc2_staging_region(self.fc2_tma_stages) else: self.fc2_tma_stages = 0 self.fc2_staging_spec = None @@ -1274,9 +1131,7 @@ def epilogue_sync_barrier(cls) -> pipeline.NamedBarrier: outside the subtile loop, FC1's amax exchange and FC2's bulk-store handshake belong to different work tiles. """ - return pipeline.NamedBarrier( - barrier_id=cls._EpilogueSyncWaitBarId, num_threads=32 * cls._EpilogueWarpCnt - ) + return pipeline.NamedBarrier(barrier_id=cls._EpilogueSyncWaitBarId, num_threads=32 * cls._EpilogueWarpCnt) def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: """Declare the epilogue scratch as one allocation shared by two exclusive lifetimes. @@ -1331,13 +1186,9 @@ def fc2_tma_staged_smem_spec(self, n_stages: int) -> Tuple[Tuple, Tuple, Tuple[i raise ValueError(f"Unsupported FC2 TMA staging dtype {wire_dtype}.") return shape, stride, swizzle - def fc2_tma_staged_smem_layout( - self, n_stages: int, without_stage_mode: bool = False - ) -> cute.ComposedLayout: + def fc2_tma_staged_smem_layout(self, n_stages: int, without_stage_mode: bool = False) -> cute.ComposedLayout: shape, stride, swizzle = self.fc2_tma_staged_smem_spec(n_stages) - layout = cute.make_composed_layout( - cute.make_swizzle(*swizzle), 0, cute.make_layout(shape, stride=stride) - ) + layout = cute.make_composed_layout(cute.make_swizzle(*swizzle), 0, cute.make_layout(shape, stride=stride)) if without_stage_mode: return cute.select(layout, mode=[0, 1]) return layout @@ -1377,10 +1228,7 @@ def _make_fc2_staging_region(self, stage_count: int) -> SmemRegion: kind="tensor", dtype=single_stage_region.dtype, shape=(*single_stage_region.shape, stage_count), - stride=( - *single_stage_region.stride, - single_stage_region.cosize if stage_count > 1 else 0, - ), + stride=(*single_stage_region.stride, single_stage_region.cosize if stage_count > 1 else 0), swizzle=single_stage_region.swizzle, byte_alignment=single_stage_region.byte_alignment, ) @@ -1454,17 +1302,12 @@ def run( fc2_smem_tensor = smem_workspace.tensor(self.fc2_staging_region, smem_base) else: fc2_smem_tensor = None - if cutlass.const_expr( - self.fc2_use_tma and (tma_atom_fc2_output is None or fc2_tma_output is None) - ): + if cutlass.const_expr(self.fc2_use_tma and (tma_atom_fc2_output is None or fc2_tma_output is None)): raise ValueError("FC2 TMA store requires a TMA atom and token-major output tensor.") if cutlass.const_expr( - self.communication_enabled - and (token_src_metadata is None or peer_rank_ptr_mapper is None) + self.communication_enabled and (token_src_metadata is None or peer_rank_ptr_mapper is None) ): - raise ValueError( - "Communication-enabled Epilogue requires token metadata and a peer pointer mapper." - ) + raise ValueError("Communication-enabled Epilogue requires token metadata and a peer pointer mapper.") if cutlass.const_expr(self.token_back_enabled and fc2_done_counter is None): raise ValueError("Token-back requires an FC2 done counter.") if cutlass.const_expr(self.token_back_push_sf and fc2_output_sf is None): @@ -1553,13 +1396,9 @@ def run( # Publish completion for the work tile snapshotted above. if cur_was_linear1: - flag_tracker = fc1_epi.signal_fc1_done( - prev_work_tile_info, work_tile_info, flag_tracker - ) + flag_tracker = fc1_epi.signal_fc1_done(prev_work_tile_info, work_tile_info, flag_tracker) else: - flag_tracker = fc2_epi.signal_fc2_done( - prev_work_tile_info, work_tile_info, flag_tracker - ) + flag_tracker = fc2_epi.signal_fc2_done(prev_work_tile_info, work_tile_info, flag_tracker) # Tail flush flag_tracker.fire() @@ -1569,9 +1408,7 @@ class _ImmutableAfterInit: def __setattr__(self, name, value): if self.__dict__.get("_frozen_", False): - raise AttributeError( - f"{type(self).__name__} is immutable after __init__ (cannot set {name!r})." - ) + raise AttributeError(f"{type(self).__name__} is immutable after __init__ (cannot set {name!r}).") object.__setattr__(self, name, value) def _freeze(self) -> None: @@ -1600,8 +1437,7 @@ def __init__( # spelled out here, so an 8-bit fc1 output picks up its own atom without a code change. staged_layout = base.fc1_staged_smem_layout(base.subtile_cnt) self.smem_tensor = cute.make_tensor( - cute.recast_ptr(staging_pointer, staged_layout.inner, dtype=base.fc1_output_dtype), - staged_layout.outer, + cute.recast_ptr(staging_pointer, staged_layout.inner, dtype=base.fc1_output_dtype), staged_layout.outer ) # Kept unswizzled and untyped so fc1_quant can carve an fp32 amax-exchange view out of one # staging stage without going through the staging tensor's swizzle. @@ -1637,19 +1473,14 @@ def signal_fc1_done(self, work_tile_info, next_work_tile_info, flag_tracker): or self.intermediate_downproj % self.cluster_tile_intermediate_downproj != 0 ) if cutlass.const_expr(needs_intermediate_guard): - in_bound = ( - work_tile_info.tile_m_idx * self._EpilogueFc1IntermediateDownTileSize - < self.fc1_output.shape[1] - ) + in_bound = work_tile_info.tile_m_idx * self._EpilogueFc1IntermediateDownTileSize < self.fc1_output.shape[1] else: in_bound = True slot = work_tile_info.cumulative_token_block_count + work_tile_info.tile_n_idx flag_address = Int64(0) if in_bound: flag_address = (self.fc1_done_counter.iterator + slot).toint() - return flag_tracker.accumulate( - next_work_tile_info.phase, self.fc1_epi_flag_batch, flag_address - ) + return flag_tracker.accumulate(next_work_tile_info.phase, self.fc1_epi_flag_batch, flag_address) @cute.jit def __call__( @@ -1661,13 +1492,9 @@ def __call__( is_odd_turn: cutlass.Int32, ): # (tokens_this_expert, intermediate_down, 1) - real_fc1_output, _ = self.kernel_extension.get_gmem_tensor( - "c", self.fc1_output, work_tile_info - ) + real_fc1_output, _ = self.kernel_extension.get_gmem_tensor("c", self.fc1_output, work_tile_info) # (tokens_this_expert, intermediate_down, 1) - real_fc1_output_sf, _ = self.kernel_extension.get_gmem_tensor( - "sfc", self.fc1_output_sf, work_tile_info - ) + real_fc1_output_sf, _ = self.kernel_extension.get_gmem_tensor("sfc", self.fc1_output_sf, work_tile_info) # subtile-irrevalent hoist out here. if cutlass.const_expr(self.optional_epi_args.fc1_alpha is not None): alpha_val = self.optional_epi_args.fc1_alpha[work_tile_info.expert_idx] @@ -1679,8 +1506,7 @@ def __call__( norm_const = None # (cta_tile_m, cta_tile_n) -> (epi_tile_m, epi_tile_n, iters) tmem_acc_tensor_tiled_by_epi_tile = cute.flat_divide( - tmem_acc_tensor, - (self._EpilogueFc1IntermediateGateUpTileSize, self._EpilogueTokenTileSize), + tmem_acc_tensor, (self._EpilogueFc1IntermediateGateUpTileSize, self._EpilogueTokenTileSize) )[None, None, 0, None] acc_pipeline.consumer_wait(acc_consumer_state) @@ -1691,12 +1517,8 @@ def __call__( remain_subtile_cnt = self.subtile_cnt - unroll_tile_cnt if cutlass.const_expr(unroll_tile_cnt > 0): - subtile_idx_first = (cutlass.Int32(self.subtile_cnt) - is_odd_turn) % cutlass.Int32( - self.subtile_cnt - ) - subtile_idx_second = ( - cutlass.Int32(self.subtile_cnt + 1) - is_odd_turn - ) % cutlass.Int32(self.subtile_cnt) + subtile_idx_first = (cutlass.Int32(self.subtile_cnt) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) + subtile_idx_second = (cutlass.Int32(self.subtile_cnt + 1) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) # preload_subtile_first: subtile_idx_first's raw PRE-transpose acc, LDTM'd by # all 128 epi threads into 4 reg tensors == the 4 quadrants of the subtile's @@ -1735,9 +1557,7 @@ def __call__( self.run_subtile( work_tile_info=work_tile_info, subtile_idx=subtile_idx_pair[i], - tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[ - None, None, subtile_idx_second - ], + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx_second], preload_acc=preload_pair[i], fc1_output=real_fc1_output, fc1_output_sf=real_fc1_output_sf, @@ -1748,9 +1568,7 @@ def __call__( for i in cutlass.range(remain_subtile_cnt, unroll=1): real_i = i + unroll_tile_cnt if cutlass.const_expr(self.overlapping_accum): - subtile_idx = ( - cutlass.Int32(real_i + self.subtile_cnt) - is_odd_turn - ) % cutlass.Int32(self.subtile_cnt) + subtile_idx = (cutlass.Int32(real_i + self.subtile_cnt) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) else: subtile_idx = cutlass.Int32(real_i) @@ -1802,9 +1620,7 @@ def run_subtile( # Step -1: preload topk scores. current_two_token_idices = ( - work_tile_info.tile_n_idx * self.cta_tile_n - + subtile_idx * self._EpilogueTokenTileSize - + self.lane_idx, + work_tile_info.tile_n_idx * self.cta_tile_n + subtile_idx * self._EpilogueTokenTileSize + self.lane_idx, work_tile_info.tile_n_idx * self.cta_tile_n + subtile_idx * self._EpilogueTokenTileSize + self.lane_idx @@ -1828,11 +1644,8 @@ def run_subtile( up_token_32_64 = cute.make_rmem_tensor((16,), cutlass.Float32) # Although hardcode is not right, but since the whole tmem transpose is too tricky, I have to hardcode... # (epi_tile_m, epi_tile_n) -> (warp_local_epi_tile_m, epi_tile_n) - # tmem_subtile_tensor_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[(None, - # self.warp_idx), None] - tmem_subtile_tensor_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[ - (None, 0), None - ] + # tmem_subtile_tensor_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[(None, self.warp_idx), None] + tmem_subtile_tensor_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[(None, 0), None] # (warp_local_epi_tile_m, epi_tile_n) -> (((16, 32), 1), (2, 2)) tmem_subtile_tensor_in_first_load_view = cute.logical_divide( cute.zipped_divide(tmem_subtile_tensor_per_warp, (16, 32)), ((16, 32), 1) @@ -1860,16 +1673,12 @@ def run_subtile( ) # Step 1: perform swiglu on the first part, interleave with the second's 32x32 tmem transpose. - token_0_32_pre_quant_pre_trans = self.alpha_swiglu_clamp( - gate_token_0_32, up_token_0_32, alpha_val - ) + token_0_32_pre_quant_pre_trans = self.alpha_swiglu_clamp(gate_token_0_32, up_token_0_32, alpha_val) # gate_token_32_64 / up_token_32_64 are already in the transpose input # distribution (see TmemTranspose16x32 / load_subtile_raw_acc). token_32_64_tmem_trans = TmemTranspose32x32Inplace( - tmem_subtile_tensor.iterator, - reg_tensor_top=gate_token_32_64, - reg_tensor_bot=up_token_32_64, + tmem_subtile_tensor.iterator, reg_tensor_top=gate_token_32_64, reg_tensor_bot=up_token_32_64 ) # Transpose output: each lane holds (token_1, intermediate_16); tmem_dp @@ -1907,8 +1716,7 @@ def run_subtile( ] # (cta_token, cta_intermediate_down) -> (token_64, intermediate_64) fc1_gmem_subtile_view = cute.flat_divide( - fc1_gmem_cta_view, - (self._EpilogueTokenTileSize, self._EpilogueFc1IntermediateDownTileSize), + fc1_gmem_cta_view, (self._EpilogueTokenTileSize, self._EpilogueFc1IntermediateDownTileSize) )[None, None, subtile_idx, 0] tma_smem_src, tma_gmem_dst = cpasync.tma_partition( self.fc1_tma_atom, @@ -1944,21 +1752,15 @@ def alpha_swiglu_clamp( # raise fires during compilation rather than emitting a runtime branch. for _name, _t in (("gate_rmem", gate_rmem), ("up_rmem", up_rmem)): if cutlass.const_expr(_t.element_type is not cutlass.Float32): - raise TypeError( - f"alpha_swiglu_clamp: {_name} must be Float32, got {_t.element_type}" - ) + raise TypeError(f"alpha_swiglu_clamp: {_name} must be Float32, got {_t.element_type}") if cutlass.const_expr(_t.memspace != AddressSpace.rmem): raise ValueError( f"alpha_swiglu_clamp: {_name} must be a register (rmem) tensor, got address space {_t.memspace}" ) if cutlass.const_expr(cute.rank(_t) != 1): - raise ValueError( - f"alpha_swiglu_clamp: {_name} must be 1D, got rank {cute.rank(_t)}" - ) + raise ValueError(f"alpha_swiglu_clamp: {_name} must be 1D, got rank {cute.rank(_t)}") if cutlass.const_expr(cute.size(_t) % 2 != 0): - raise ValueError( - f"alpha_swiglu_clamp: {_name} element count must be even, got {cute.size(_t)}" - ) + raise ValueError(f"alpha_swiglu_clamp: {_name} element count must be even, got {cute.size(_t)}") if cutlass.const_expr(cute.size(gate_rmem) != cute.size(up_rmem)): raise ValueError( "alpha_swiglu_clamp: gate_rmem and up_rmem must have equal size, got " @@ -1989,10 +1791,7 @@ def alpha_swiglu_clamp( activation_alpha = 1.0 if self.swiglu_alpha is None else self.swiglu_alpha neg_activation_alpha_log2e = -activation_alpha * log2_e - neg_log2e_pair = ( - cutlass.Float32(neg_activation_alpha_log2e), - cutlass.Float32(neg_activation_alpha_log2e), - ) + neg_log2e_pair = (cutlass.Float32(neg_activation_alpha_log2e), cutlass.Float32(neg_activation_alpha_log2e)) one_pair = (cutlass.Float32(1.0), cutlass.Float32(1.0)) if cutlass.const_expr(self.swiglu_beta is not None): beta_pair = (cutlass.Float32(self.swiglu_beta), cutlass.Float32(self.swiglu_beta)) @@ -2071,6 +1870,7 @@ def _sigmoid(p0, p1): sigmoid_pair = _sigmoid(g0, g1) if cutlass.const_expr(self.situ_beta is None): + # SwiGLU: out = (up + swiglu_beta) * gate * sigmoid(swiglu_alpha * gate) # The SwiGLUBias offset follows the up clamp. if cutlass.const_expr(self.swiglu_beta is not None): u0, u1 = cute.arch.add_packed_f32x2((u0, u1), beta_pair) @@ -2116,9 +1916,7 @@ def _sigmoid(p0, p1): def fc1_quant( self, work_tile_info: SwapAbFc12WorkTileInfo, - two_token: Tuple[ - cute.Tensor, cute.Tensor - ], # two rmem tensor, each fp32 @ (token_1, intermediate_16) + two_token: Tuple[cute.Tensor, cute.Tensor], # two rmem tensor, each fp32 @ (token_1, intermediate_16) topk_scores: Optional[Tuple[cutlass.Float32, cutlass.Float32]], norm_const: Optional[cutlass.Float32], intermediate_output_size: cutlass.Int32, @@ -2152,8 +1950,7 @@ def fc1_quant( cute.recast_ptr(self.staging_pointer, dtype=cutlass.Float32) + subtile_idx * cutlass.Int32(self.fc1_staging_stage_bytes // 4), cute.make_layout( - (self.fc1_amax_token_chunks, self._EpilogueWarpCnt * 32), - stride=(1, self.fc1_amax_token_chunks), + (self.fc1_amax_token_chunks, self._EpilogueWarpCnt * 32), stride=(1, self.fc1_amax_token_chunks) ), ), 16, @@ -2180,13 +1977,8 @@ def fc1_quant( stores_scale_factor = self.warp_idx % 2 == 0 else: stores_scale_factor = True - subtile_token_start = ( - work_tile_info.tile_n_idx * self.cta_tile_n + subtile_idx * self._EpilogueTokenTileSize - ) - token_idx_pair = ( - subtile_token_start + self.lane_idx, - subtile_token_start + self.lane_idx + 32, - ) + subtile_token_start = work_tile_info.tile_n_idx * self.cta_tile_n + subtile_idx * self._EpilogueTokenTileSize + token_idx_pair = (subtile_token_start + self.lane_idx, subtile_token_start + self.lane_idx + 32) # This subtile's (token, intermediate) shared output stage, tiled into (1, 16) blocks so # each thread's cells slice out directly (zipped_divide + slice; avoids the ambiguous @@ -2243,9 +2035,7 @@ def fc1_quant( # 4) STS the quantized values into this subtile's shared output stage. # ((1, 16), (token_tile_size, warp_cnt)) -> (16) smem_thread_row = smem_tiled[(0, None), (self.lane_idx + 32 * half, self.warp_idx)] - cute.copy( - store_atom, cute.coalesce(data_by_token[None, half]), cute.coalesce(smem_thread_row) - ) + cute.copy(store_atom, cute.coalesce(data_by_token[None, half]), cute.coalesce(smem_thread_row)) @dataclasses.dataclass(frozen=True) @@ -2343,20 +2133,15 @@ def signal_fc2_done(self, work_tile_info, next_work_tile_info, flag_tracker): if cutlass.const_expr(publish): flag_address = (self.fc2_done_counter.iterator + work_tile_info.expert_idx).toint() no_fire: cutlass.Constexpr = not publish - return flag_tracker.accumulate( - next_work_tile_info.phase, self.fc2_epi_flag_batch, flag_address, no_fire - ) + return flag_tracker.accumulate(next_work_tile_info.phase, self.fc2_epi_flag_batch, flag_address, no_fire) @cute.jit def _make_output_router(self, work_tile_info: SwapAbFc12WorkTileInfo) -> "Fc2OutputRouter": task_tile_data_row_start = ( - work_tile_info.cumulative_data_physical_row - + work_tile_info.tile_n_idx * cutlass.Int32(self.cta_tile_n) + work_tile_info.cumulative_data_physical_row + work_tile_info.tile_n_idx * cutlass.Int32(self.cta_tile_n) ) hidden_base_this_cta_tile = work_tile_info.tile_m_idx * cutlass.Int32(self.cta_tile_m) - valid_hidden_this_cta_tile = ( - cutlass.Int32(self.fc2_output.shape[2]) - hidden_base_this_cta_tile - ) + valid_hidden_this_cta_tile = cutlass.Int32(self.fc2_output.shape[2]) - hidden_base_this_cta_tile if valid_hidden_this_cta_tile < 0: valid_hidden_this_cta_tile = 0 if valid_hidden_this_cta_tile > self._EpilogueFc2HiddenTileSize: @@ -2365,9 +2150,7 @@ def _make_output_router(self, work_tile_info: SwapAbFc12WorkTileInfo) -> "Fc2Out metadata = None peer_rank_ptr_mapper = None data_token_base = task_tile_data_row_start - if cutlass.const_expr( - self.token_src_metadata is not None and not self.token_back_push_data - ): + if cutlass.const_expr(self.token_src_metadata is not None and not self.token_back_push_data): metadata = cute.domain_offset((task_tile_data_row_start,), self.token_src_metadata) peer_rank_ptr_mapper = self.peer_rank_ptr_mapper data_token_base = None @@ -2375,10 +2158,7 @@ def _make_output_router(self, work_tile_info: SwapAbFc12WorkTileInfo) -> "Fc2Out if cutlass.const_expr(self.combine_format.is_quantized): base_outputs = (self.fc2_output, self.fc2_output_sf) token_bases = (data_token_base, task_tile_data_row_start) - output_mappings = ( - self.process_pipeline.store_out_mapping, - self.process_pipeline.sf_store_out_mapping, - ) + output_mappings = (self.process_pipeline.store_out_mapping, self.process_pipeline.sf_store_out_mapping) else: base_outputs = self.fc2_output token_bases = data_token_base @@ -2427,21 +2207,13 @@ def __call__( # Overlap path preloads two subtiles before releasing acc TMEM. unroll_tile_cnt = ( - 2 - if cutlass.const_expr( - self.overlapping_accum and self.process_pipeline.require_tmem_trans - ) - else 0 + 2 if cutlass.const_expr(self.overlapping_accum and self.process_pipeline.require_tmem_trans) else 0 ) remain_subtile_cnt = self.subtile_cnt - unroll_tile_cnt if cutlass.const_expr(unroll_tile_cnt > 0): - subtile_idx_first = (cutlass.Int32(self.subtile_cnt) - is_odd_turn) % cutlass.Int32( - self.subtile_cnt - ) - subtile_idx_second = ( - cutlass.Int32(self.subtile_cnt + 1) - is_odd_turn - ) % cutlass.Int32(self.subtile_cnt) + subtile_idx_first = (cutlass.Int32(self.subtile_cnt) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) + subtile_idx_second = (cutlass.Int32(self.subtile_cnt + 1) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) # Each warp preloads its local 32-hidden x 64-token accumulator as two # (32,) FP32 tensors: @@ -2453,10 +2225,7 @@ def __call__( # consecutive register pair holds adjacent token columns at one # hidden coordinate. After F2FP, that pair is one BF16x2 MOVM input. preload_subtile_first: Tuple[cute.Tensor, ...] = self.process_pipeline.tmem_acc_load( - tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[ - None, None, subtile_idx_first - ], - epi=self, + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx_first], epi=self ) # Release acc to next MMA unconditionally. @@ -2466,10 +2235,7 @@ def __call__( # Same two-half mapping for the other token subtile. Both preloaded # subtiles can now use this second subtile's TMEM as transpose workspace. preload_subtile_second: Tuple[cute.Tensor, ...] = self.process_pipeline.tmem_acc_load( - tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[ - None, None, subtile_idx_second - ], - epi=self, + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx_second], epi=self ) # Both unrolled subtiles borrow tmem_subtile_second as workspace. @@ -2481,9 +2247,7 @@ def __call__( work_tile_info=work_tile_info, epilogue_iter_idx=cutlass.Int32(i), subtile_idx=subtile_idx_pair[i], - tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[ - None, None, subtile_idx_second - ], + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx_second], preload_acc=preload_pair[i], fc2_output_router=fc2_output_router, alpha_val=alpha_val, @@ -2500,9 +2264,7 @@ def __call__( # for i in cutlass.range_constexpr(remain_subtile_cnt): real_i = i + unroll_tile_cnt if cutlass.const_expr(self.overlapping_accum): - subtile_idx = ( - cutlass.Int32(real_i + self.subtile_cnt) - is_odd_turn - ) % cutlass.Int32(self.subtile_cnt) + subtile_idx = (cutlass.Int32(real_i + self.subtile_cnt) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) else: subtile_idx = cutlass.Int32(real_i) @@ -2543,9 +2305,7 @@ def run_subtile( ): process_pipeline = self.process_pipeline if cutlass.const_expr(preload_acc is None): - loaded = process_pipeline.tmem_acc_load( - tmem_subtile_tensor=tmem_subtile_tensor, epi=self - ) + loaded = process_pipeline.tmem_acc_load(tmem_subtile_tensor=tmem_subtile_tensor, epi=self) if release_after_ldtm: cute.arch.fence_view_async_tmem_load() acc_pipeline.consumer_release(acc_consumer_state) @@ -2555,9 +2315,7 @@ def run_subtile( casted = process_pipeline.f2fp(*loaded, alpha_val=alpha_val) # reorder returns a bare RMEM fragment in the store's expected pre-store # distribution; reorder + store are paired 1:1 inside the pipeline. - pre_store = process_pipeline.post_f2fp_reorder( - casted=casted, tmem_subtile_view=tmem_subtile_tensor - ) + pre_store = process_pipeline.post_f2fp_reorder(casted=casted, tmem_subtile_view=tmem_subtile_tensor) process_pipeline.store_function( epi=self, subtile=pre_store, @@ -2586,9 +2344,7 @@ class Fc2OutputRouter: epi_tid: cutlass.Int32 combine_format: CombineFormat # After metadata prefetch - dst_ptrs: Optional[cute.Tensor] = ( - None # i64 x (copy_iters_this_thread_cta_tile), fundamentally the pointers. - ) + dst_ptrs: Optional[cute.Tensor] = None # i64 x (copy_iters_this_thread_cta_tile), fundamentally the pointers. valid: Optional[cute.Tensor] = None # (copy_iters_this_thread_cta_tile) @property @@ -2610,11 +2366,7 @@ def sf_token_base(self) -> Optional[cutlass.Int32]: @property def data_mapping(self) -> FunctionMapping: - return ( - self.output_mappings[0] - if isinstance(self.output_mappings, tuple) - else self.output_mappings - ) + return self.output_mappings[0] if isinstance(self.output_mappings, tuple) else self.output_mappings @property def sf_mapping(self) -> Optional[FunctionMapping]: @@ -2622,9 +2374,7 @@ def sf_mapping(self) -> Optional[FunctionMapping]: def __post_init__(self) -> None: if (self.metadata is None) == (self.data_token_base is None): - raise ValueError( - "Fc2OutputRouter requires exactly one of metadata or a (data) token base." - ) + raise ValueError("Fc2OutputRouter requires exactly one of metadata or a (data) token base.") if (self.metadata is None) != (self.peer_rank_ptr_mapper is None): raise ValueError("Fc2OutputRouter requires peer_rank_ptr_mapper iff metadata is set.") if self.reduce_topk_in_epilogue and self.metadata is None: @@ -2643,8 +2393,7 @@ def prefetch(self) -> "Fc2OutputRouter": valid = cute.make_rmem_tensor((copy_iters,), cutlass.Int32) dst_ptrs = cute.make_rmem_tensor((copy_iters,), cutlass.Int64) - # Compiler should be able to optimize the same token_copy_group's offset add. (Fundamental cse + - # strength_reduce) + # Compiler should be able to optimize the same token_copy_group's offset add. (Fundamental cse + strength_reduce) # We should check the SASS to ensure this happens. for iter_idx in cutlass.range_constexpr(copy_iters): coord = self.data_mapping.evaluate(epi_tid=self.epi_tid, iter_idx=iter_idx) @@ -2663,14 +2412,11 @@ def prefetch(self) -> "Fc2OutputRouter": dst_hidden = hidden_in_tile + self.hidden_base_this_cta_tile # Int64 token coord: dst_tokens*K*H overflows int32 once # T*K*H exceeds 2^31 (data_output is (token, topk, hidden)). - dst_ptrs[iter_idx] = self.data_output[ - Int64(dst_tokens), None, dst_hidden - ].iterator.toint() + dst_ptrs[iter_idx] = self.data_output[Int64(dst_tokens), None, dst_hidden].iterator.toint() else: md = TokenSrcMetadata.load( - self.metadata.iterator.toint() - + Int64(token_in_tile) * Int64(TokenSrcMetadata.nbytes) + self.metadata.iterator.toint() + Int64(token_in_tile) * Int64(TokenSrcMetadata.nbytes) ) dst_rank = md.src_rank dst_token = md.src_token @@ -2682,9 +2428,7 @@ def prefetch(self) -> "Fc2OutputRouter": # Int64 token coord: domain_offset on (token, topk, hidden) # computes dst_token*K*H, which overflows int32 once T*K*H > 2^31. dst_ptrs[iter_idx] = self.peer_rank_ptr_mapper.map_pointer( - cute.domain_offset( - (Int64(dst_token), dst_topk, dst_hidden), self.data_output - ).iterator, + cute.domain_offset((Int64(dst_token), dst_topk, dst_hidden), self.data_output).iterator, dst_rank, byte_alignment=32, ).toint() @@ -2692,9 +2436,7 @@ def prefetch(self) -> "Fc2OutputRouter": return dataclasses.replace(self, dst_ptrs=dst_ptrs, valid=valid) @cute.jit - def get_data_dst( - self, iter_idx: Union[int, cutlass.Int32] - ) -> Tuple[cute.Pointer, cutlass.Int32]: + def get_data_dst(self, iter_idx: Union[int, cutlass.Int32]) -> Tuple[cute.Pointer, cutlass.Int32]: """Per-issue DATA destination: gmem pointer + validity predicate. The router owns ``data_output`` so the caller never re-assembles a @@ -2726,9 +2468,7 @@ def get_data_dst( # comm: read the pointer / validity prefetched by prefetch(). addr = self.dst_ptrs[iter_idx] pred = self.valid[iter_idx] - ptr = cute.make_ptr( - self.data_output.element_type, addr, AddressSpace.gmem, assumed_align=32 - ) + ptr = cute.make_ptr(self.data_output.element_type, addr, AddressSpace.gmem, assumed_align=32) return ptr, pred @cute.jit @@ -2755,9 +2495,7 @@ def get_sf_dst(self, iter_idx: Union[int, cutlass.Int32]) -> Tuple[cute.Pointer, addr = self.sf_output[Int64(sf_row), None, sf_hidden].iterator.toint() # Per-block scale offsets are element-granular; claim the scale dtype's # natural element alignment (e8m0 1 B / bf16 2 B). - sf_ptr = cute.make_ptr( - self.sf_output.element_type, addr, AddressSpace.gmem, assumed_align=4 - ) + sf_ptr = cute.make_ptr(self.sf_output.element_type, addr, AddressSpace.gmem, assumed_align=4) return sf_ptr, pred @@ -2781,16 +2519,13 @@ def make_fc2_stg_cta_store_out_mapping( }, ) store_out_mapping = FunctionMapping( - domain=CoordinateSpace( - ("epi_tid", "iter_idx"), (128, stgs_per_hidden32 * cta_token_tile_size // 32) - ), + domain=CoordinateSpace(("epi_tid", "iter_idx"), (128, stgs_per_hidden32 * cta_token_tile_size // 32)), codomain=CoordinateSpace( ("token_in_cta_tile", "hidden_in_cta_tile"), (cta_token_tile_size, cta_hidden_tile_size) ), function=lambda epi_tid, iter_idx: { "token_in_cta_tile": epi_tid % 32 + iter_idx // stgs_per_hidden32 * 32, - "hidden_in_cta_tile": (iter_idx % stgs_per_hidden32) * elems_per_stg - + epi_tid // 32 * 32, + "hidden_in_cta_tile": (iter_idx % stgs_per_hidden32) * elems_per_stg + epi_tid // 32 * 32, }, ) sf_store_out_mapping = None @@ -2804,8 +2539,7 @@ def stg_sf_mapping(epi_tid, iter_idx): sf_store_out_mapping = FunctionMapping( domain=CoordinateSpace(("epi_tid", "iter_idx"), (128, cta_token_tile_size // 32)), codomain=CoordinateSpace( - ("token_in_cta_tile", "hidden_in_cta_tile"), - (cta_token_tile_size, cta_hidden_tile_size), + ("token_in_cta_tile", "hidden_in_cta_tile"), (cta_token_tile_size, cta_hidden_tile_size) ), function=stg_sf_mapping, ) @@ -2835,10 +2569,7 @@ def make_fc2_redg_cta_store_out_mapping( + (epi_tid % 32) // 4 ), "hidden_in_cta_tile": ( - (epi_tid // 32) * 32 - + (epi_tid % 4) * 4 - + (((elem_idx // 4) % 4) // 2) * 16 - + elem_idx % 4 + (epi_tid // 32) * 32 + (epi_tid % 4) * 4 + (((elem_idx // 4) % 4) // 2) * 16 + elem_idx % 4 ), }, ) @@ -2858,17 +2589,13 @@ def make_fc2_redg_cta_store_out_mapping( + ((iter_idx % 4) % 2) * 8 + (epi_tid % 32) // 4 ), - "hidden_in_cta_tile": ( - (epi_tid // 32) * 32 + (epi_tid % 4) * 4 + ((iter_idx % 4) // 2) * 16 - ), + "hidden_in_cta_tile": ((epi_tid // 32) * 32 + (epi_tid % 4) * 4 + ((iter_idx % 4) // 2) * 16), }, ) return store_out_mapping, None, fundamental_mapping -def make_fc2_ublk_store_out_mapping( - combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int -): +def make_fc2_ublk_store_out_mapping(combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int): assert cta_hidden_tile_size == 128 assert cta_token_tile_size % 64 == 0 # UBLK pushes whole hidden rows by byte count, so the token/hidden mapping @@ -2897,9 +2624,7 @@ def make_fc2_ublk_store_out_mapping( ("token_in_cta_tile", "hidden_in_cta_tile"), (max_token_cta_tile, cta_hidden_tile_size) ), function=lambda epi_tid, iter_idx: { - "token_in_cta_tile": ( - iter_idx * 128 + ((epi_tid % 32) // 16) * 64 + (epi_tid // 32) * 16 + epi_tid % 16 - ), + "token_in_cta_tile": (iter_idx * 128 + ((epi_tid % 32) // 16) * 64 + (epi_tid // 32) * 16 + epi_tid % 16), "hidden_in_cta_tile": 0, }, ) @@ -2933,9 +2658,7 @@ def fc2_f2fp(*tensors, alpha_val: Optional[cutlass.Float32] = None, **_) -> cute elems_processed = 0 for t in tensors: current_tensor_size = cute.size(t) - dst = cute.make_tensor( - converted_acc.iterator + elems_processed, cute.make_layout((current_tensor_size,)) - ) + dst = cute.make_tensor(converted_acc.iterator + elems_processed, cute.make_layout((current_tensor_size,))) if cutlass.const_expr(alpha_val is None): dst.store(t.load().to(reorder_dtype)) else: @@ -2965,9 +2688,7 @@ def fc2_stg_tmem_acc_load(*, tmem_subtile_tensor: cute.Tensor, **_): ptr = tmem_subtile_tensor.iterator half_lane_offset = 16 * TmemTranspose32x64B16Movm._tmem_row_stride top_view = cute.make_tensor(ptr, TmemTranspose32x64B16Movm._tmem_layout(16, 64)) - bottom_view = cute.make_tensor( - ptr + half_lane_offset, TmemTranspose32x64B16Movm._tmem_layout(16, 64) - ) + bottom_view = cute.make_tensor(ptr + half_lane_offset, TmemTranspose32x64B16Movm._tmem_layout(16, 64)) top = cute.make_rmem_tensor((32,), cutlass.Float32) bottom = cute.make_rmem_tensor((32,), cutlass.Float32) cute.copy(atom_ld16x256, top_view, TmemTranspose32x64B16Movm._rmem_copy_view(top, 32)) @@ -2980,13 +2701,9 @@ def fc2_ublk_tmem_acc_load(*, tmem_subtile_tensor: cute.Tensor, epi, **_): # UBLK consumes a warp-local 32-hidden x 64-token slice. The caller passes # the CTA-level 128-hidden x 64-token subtile view, so select this epi # warp's hidden block before issuing LDTM.x64. - tmem_subtile_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[ - (None, epi.warp_idx), None - ] + tmem_subtile_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[(None, epi.warp_idx), None] raw_regs = cute.make_rmem_tensor((64,), cutlass.Float32) - atom_ld32x32_x64 = cute.make_copy_atom( - tcgen05.Ld32x32bOp(tcgen05.Repetition.x64), cutlass.Float32 - ) + atom_ld32x32_x64 = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x64), cutlass.Float32) cute.copy( atom_ld32x32_x64, wrap_into_copy_standard_layout(tmem_subtile_per_warp), @@ -3011,20 +2728,14 @@ def fc2_stg_post_f2fp_reorder( def fc2_redg_post_f2fp_reorder(*, casted: cute.Tensor, tmem_subtile_view: cute.Tensor, **_): # (epi_tid, elem_idx) -> (token_64, hidden_128), each thread hold token_2 x hidden_32 natural = fc2_stg_post_f2fp_reorder(casted=casted, tmem_subtile_view=tmem_subtile_view) - core_matrix_reorder_sttm_atom = cute.make_copy_atom( - tcgen05.St32x32bOp(tcgen05.Repetition.x16), cutlass.Float32 - ) - core_matrix_reorder_ldtm_atom = cute.make_copy_atom( - tcgen05.Ld16x256bOp(tcgen05.Repetition.x2), cutlass.Float32 - ) + core_matrix_reorder_sttm_atom = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x16), cutlass.Float32) + core_matrix_reorder_ldtm_atom = cute.make_copy_atom(tcgen05.Ld16x256bOp(tcgen05.Repetition.x2), cutlass.Float32) # ((16, 2), token_32_group) token_groups = cute.logical_divide(cute.zipped_divide(natural, (32,)), (16, None)) out = cute.make_rmem_tensor(token_groups.shape, casted.dtype) out_as_i32 = cute.recast_tensor(out, cutlass.Float32) # (32, 64) - tmem_warp = cute.flat_divide(tmem_subtile_view, (32, cute.size(tmem_subtile_view, 1)))[ - None, None, 0, 0 - ] + tmem_warp = cute.flat_divide(tmem_subtile_view, (32, cute.size(tmem_subtile_view, 1)))[None, None, 0, 0] # (16, 16, 16dp_group, token_32_groups). Note, this tmem can provide 2x cols since the original is bf16. tmem_groups = cute.flat_divide(tmem_warp, (16, 16)) for group_idx in cutlass.range_constexpr(cute.size(token_groups, 1)): @@ -3064,9 +2775,7 @@ def fc2_stg_store_function( sf_regs = None stg_width_elems: cutlass.Constexpr[int] = min(32, 256 // data_subtile.element_type.width) stg_bits: cutlass.Constexpr[int] = stg_width_elems * data_subtile.element_type.width - copy_atom_vec = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=stg_bits - ) + copy_atom_vec = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=stg_bits) elems_per_thread: cutlass.Constexpr[int] = cute.size(data_subtile) if cutlass.const_expr(elems_per_thread % stg_width_elems != 0): raise ValueError( @@ -3078,9 +2787,7 @@ def fc2_stg_store_function( sf_scales_per_stg: cutlass.Constexpr[int] = 32 // epi.combine_format.scale_block token_groups_per_subtile: cutlass.Constexpr[int] = epi._EpilogueTokenTileSize // 32 for token_group in cutlass.range_constexpr(token_groups_per_subtile): - sf_iter = ( - cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group - ) + sf_iter = cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group sf_ptr, sf_pred = fc2_output_router.get_sf_dst(sf_iter) if sf_pred != cutlass.Int32(0): sf_dst = cute.make_tensor(sf_ptr, cute.make_layout((sf_scales_per_stg,))) @@ -3097,11 +2804,7 @@ def fc2_stg_store_function( if pred != cutlass.Int32(0): src_i = cute.make_tensor(copy_src[None, local_iter].iterator, single_copy_layout) dst_i = cute.make_tensor(dst_ptr, single_copy_layout) - cute.copy( - copy_atom_vec, - cute.recast_tensor(src_i, cutlass.Int32), - cute.recast_tensor(dst_i, cutlass.Int32), - ) + cute.copy(copy_atom_vec, cute.recast_tensor(src_i, cutlass.Int32), cute.recast_tensor(dst_i, cutlass.Int32)) @cute.jit @@ -3115,12 +2818,8 @@ def fc2_tma_store_function( fc2_output_router: Fc2OutputRouter, **_, ): - if cutlass.const_expr( - epi.smem_tensor is None or epi.tma_atom_fc2_output is None or epi.fc2_tma_output is None - ): - raise ValueError( - "FC2 TMA store requires staged SMEM, a TMA atom, and a token-major output tensor." - ) + if cutlass.const_expr(epi.smem_tensor is None or epi.tma_atom_fc2_output is None or epi.fc2_tma_output is None): + raise ValueError("FC2 TMA store requires staged SMEM, a TMA atom, and a token-major output tensor.") if cutlass.const_expr(epi.combine_format.is_quantized): data_subtile, sf_regs = QuantImpl(epi.combine_format, "regs_in_thread")(subtile) @@ -3132,9 +2831,7 @@ def fc2_tma_store_function( sf_scales_per_stg: cutlass.Constexpr[int] = 32 // epi.combine_format.scale_block token_groups_per_subtile: cutlass.Constexpr[int] = epi._EpilogueTokenTileSize // 32 for token_group in cutlass.range_constexpr(token_groups_per_subtile): - sf_iter = ( - cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group - ) + sf_iter = cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group sf_ptr, sf_pred = fc2_output_router.get_sf_dst(sf_iter) if sf_pred != cutlass.Int32(0): sf_dst = cute.make_tensor(sf_ptr, cute.make_layout((sf_scales_per_stg,))) @@ -3148,15 +2845,11 @@ def fc2_tma_store_function( or data_subtile.stride[0] != 1 or 32 % vector_elements != 0 ): - raise ValueError( - "FC2 TMA store requires a contiguous two-token by 32-hidden register fragment." - ) + raise ValueError("FC2 TMA store requires a contiguous two-token by 32-hidden register fragment.") stage_idx = epilogue_iter_idx % cutlass.Int32(epi.fc2_tma_stages) smem_stage = epi.smem_tensor[None, None, stage_idx] - store_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), data_subtile.element_type, num_bits_per_copy=128 - ) + store_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), data_subtile.element_type, num_bits_per_copy=128) tiler_mn = (64, 128) layout_copy_tv = cute.make_layout(((32, 4), (32, 2)), stride=((1, 2048), (64, 32))) tiled_store = cute.make_tiled_copy(store_atom, layout_copy_tv, tiler_mn) @@ -3170,9 +2863,7 @@ def fc2_tma_store_function( + work_tile_info.tile_n_idx * cutlass.Int32(epi.cta_tile_n) + subtile_idx * cutlass.Int32(epi._EpilogueTokenTileSize) ) // cutlass.Int32(epi._EpilogueTokenTileSize) - tiled_output = cute.flat_divide( - epi.fc2_tma_output, (epi._EpilogueTokenTileSize, epi._EpilogueFc2HiddenTileSize) - ) + tiled_output = cute.flat_divide(epi.fc2_tma_output, (epi._EpilogueTokenTileSize, epi._EpilogueFc2HiddenTileSize)) gmem_subtile = tiled_output[None, None, token_tile_idx, work_tile_info.tile_m_idx, 0] tma_smem_src, tma_gmem_dst = cpasync.tma_partition( epi.tma_atom_fc2_output, @@ -3212,9 +2903,7 @@ def fc2_ublk_store_function_impl( sf_scales_per_stg: cutlass.Constexpr[int] = 32 // epi.combine_format.scale_block token_groups_per_subtile: cutlass.Constexpr[int] = epi._EpilogueTokenTileSize // 32 for token_group in cutlass.range_constexpr(token_groups_per_subtile): - sf_iter = ( - cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group - ) + sf_iter = cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group sf_ptr, sf_pred = fc2_output_router.get_sf_dst(sf_iter) if sf_pred != cutlass.Int32(0): sf_dst = cute.make_tensor(sf_ptr, cute.make_layout((sf_scales_per_stg,))) @@ -3241,12 +2930,8 @@ def fc2_ublk_store_function_impl( or data_subtile.stride[0] != 1 or 32 % vector_elements != 0 ): - raise ValueError( - "FC2 UBLK store requires a contiguous two-token by 32-hidden register fragment." - ) - store_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), data_subtile.element_type, num_bits_per_copy=128 - ) + raise ValueError("FC2 UBLK store requires a contiguous two-token by 32-hidden register fragment.") + store_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), data_subtile.element_type, num_bits_per_copy=128) tiler_mn = (64, 128) layout_copy_tv = cute.make_layout(((32, 4), (32, 2)), stride=((1, 2048), (64, 32))) tiled_store = cute.make_tiled_copy(store_atom, layout_copy_tv, tiler_mn) @@ -3256,9 +2941,7 @@ def fc2_ublk_store_function_impl( cute.copy(tiled_store, register_partition, smem_partition) else: if cutlass.const_expr(cute.size(data_subtile) != epi._EpilogueTokenTileSize): - raise ValueError( - "BF16 UBLK staging requires one register per token in the 64-token subtile." - ) + raise ValueError("BF16 UBLK staging requires one register per token in the 64-token subtile.") warp_hidden_base = cutlass.Int32(warp_idx * 32) for token_idx in cutlass.range_constexpr(epi._EpilogueTokenTileSize): smem_stage[token_idx, warp_hidden_base + lane_idx] = data_subtile[token_idx] @@ -3312,18 +2995,14 @@ def fc2_redg_store_function( if pred != cutlass.Int32(0): bf16x4 = subtile_by_redg_issue[None, local_iter] packed_bf16x2 = cute.recast_tensor(bf16x4, cutlass.Float32) - red_add_relaxed_sys_v2_bf16x2( - dst_ptr, cutlass.Float32(packed_bf16x2[0]), cutlass.Float32(packed_bf16x2[1]) - ) + red_add_relaxed_sys_v2_bf16x2(dst_ptr, cutlass.Float32(packed_bf16x2[0]), cutlass.Float32(packed_bf16x2[1])) def make_fc2_stg_process_pipeline( *, combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int ) -> Fc2ProcessPipeline: - store_out_mapping, sf_store_out_mapping, fundamental_mapping = ( - make_fc2_stg_cta_store_out_mapping( - combine_format, cta_token_tile_size, cta_hidden_tile_size - ) + store_out_mapping, sf_store_out_mapping, fundamental_mapping = make_fc2_stg_cta_store_out_mapping( + combine_format, cta_token_tile_size, cta_hidden_tile_size ) return Fc2ProcessPipeline( tmem_acc_load=fc2_stg_tmem_acc_load, @@ -3340,10 +3019,8 @@ def make_fc2_stg_process_pipeline( def make_fc2_tma_process_pipeline( *, combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int ) -> Fc2ProcessPipeline: - store_out_mapping, sf_store_out_mapping, fundamental_mapping = ( - make_fc2_stg_cta_store_out_mapping( - combine_format, cta_token_tile_size, cta_hidden_tile_size - ) + store_out_mapping, sf_store_out_mapping, fundamental_mapping = make_fc2_stg_cta_store_out_mapping( + combine_format, cta_token_tile_size, cta_hidden_tile_size ) return Fc2ProcessPipeline( tmem_acc_load=fc2_stg_tmem_acc_load, @@ -3360,10 +3037,8 @@ def make_fc2_tma_process_pipeline( def make_fc2_redg_process_pipeline( *, combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int ) -> Fc2ProcessPipeline: - store_out_mapping, sf_store_out_mapping, fundamental_mapping = ( - make_fc2_redg_cta_store_out_mapping( - combine_format, cta_token_tile_size, cta_hidden_tile_size - ) + store_out_mapping, sf_store_out_mapping, fundamental_mapping = make_fc2_redg_cta_store_out_mapping( + combine_format, cta_token_tile_size, cta_hidden_tile_size ) return Fc2ProcessPipeline( tmem_acc_load=fc2_stg_tmem_acc_load, diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py index b9318a28b60f..fc47fe0e3fa5 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Block-scaled SwapAb adapter between FC12 scheduling and kernel tensor views.""" import dataclasses @@ -16,6 +13,7 @@ from .....helpers.dsl_helpers import spin_peek, spin_wait from ....schedulers.fc12_mapping import BlockPhase, SwapAbFc12WorkTileInfo, peek_ready_bit + TensorRole = Literal["a", "b", "sfa", "sfb", "c", "sfc", "topk"] @@ -54,18 +52,14 @@ def __new_from_mlir_values__(self, values: List[ir.Value]) -> "BlockScaledSwapAb def rebuild(field): nonlocal value_index field_value_count = len(extract_mlir_values(field)) - result = new_from_mlir_values( - field, values[value_index : value_index + field_value_count] - ) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) value_index += field_value_count return result fc1_done_counter_pointer = rebuild(self.fc1_done_counter_pointer) fc2_spin_threshold = rebuild(self.fc2_spin_threshold) fc1_ready_counter_pointer = ( - rebuild(self.fc1_ready_counter_pointer) - if self.fc1_ready_counter_pointer is not None - else None + rebuild(self.fc1_ready_counter_pointer) if self.fc1_ready_counter_pointer is not None else None ) if value_index != len(values): raise ValueError( @@ -91,9 +85,7 @@ def prepare_work_tile(self, work_tile: SwapAbFc12WorkTileInfo) -> SwapAbFc12Work if is_fc1: counter_pointer = self.fc1_ready_counter_pointer + counter_slot peek_flag = Int32(0) - if spin_peek( - counter_pointer, lambda value: value >= work_tile.valid_tokens_in_cta_tile - ): + if spin_peek(counter_pointer, lambda value: value >= work_tile.valid_tokens_in_cta_tile): peek_flag = Int32(peek_ready_bit) phase_and_flags = work_tile.phase_and_flags | peek_flag @@ -151,19 +143,13 @@ def get_gmem_tensor( result = cute.domain_offset((0, 0, expert_idx), tensor) per_expert_shape = (shape[0], shape[1], singleton) sf_layout = tile_atom_to_shape_SF(per_expert_shape, self.sf_vec_size) - return ( - cute.make_tensor(result.iterator, cute.make_layout(sf_layout.shape, stride=stride)), - None, - ) + return (cute.make_tensor(result.iterator, cute.make_layout(sf_layout.shape, stride=stride)), None) if cutlass.const_expr(tensor_name in ("sfb", "sfc")): result = cute.domain_offset((sf_token_offset, 0, 0), tensor) per_expert_shape = (shape[0], shape[1], singleton) sf_layout = tile_atom_to_shape_SF(per_expert_shape, self.sf_vec_size) - return ( - cute.make_tensor(result.iterator, cute.make_layout(sf_layout.shape, stride=stride)), - None, - ) + return (cute.make_tensor(result.iterator, cute.make_layout(sf_layout.shape, stride=stride)), None) if cutlass.const_expr(tensor_name == "c"): result = cute.domain_offset((data_token_offset, 0, 0), tensor) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_kernel.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_kernel.py new file mode 100644 index 000000000000..c60f1e814697 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_kernel.py @@ -0,0 +1,689 @@ +"""Lean block-scaled swap-AB FC12 kernel composition for Blackwell.""" + +from typing import ClassVar, Optional, Tuple + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +from cutlass.cute.nvgpu import OperandMajorMode +from cutlass.cutlass_dsl import Boolean +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + +from .....api import ImplDesc, KernelClass, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType +from .....helpers.cute_py_helpers import Tcgen05MmaInstruction, make_tcgen05_tmem_plan, tcgen05_block_scaled_acc_dtype +from .....helpers.device_workspace import DeviceWorkspace +from .....helpers.smem_workspace import SmemWorkspace +from .....helpers.utils import ceil_div, round_up +from .....quant_def import CombineFormat, QuantKind +from ....schedulers.base import WorkIdAcquisitionMode +from ....schedulers.fc12_scheduler import BlackwellFusedFc12Scheduler +from ...custom_mix_cga_helpers import TmaAtomOrPair +from .block_scaled_swap_ab_fc12_epilogue import GatedActEpilogueArgs, SwapABGatedActEpilogue +from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension +from .block_scaled_swap_ab_fc12_mainloop import BlockScaledSwapAbFc12Mainloop + + +class BlockScaledSwapAbFc12Kernel(KernelClass): + """Lean FC12-only top-level with one explicit eight-warp topology.""" + + fc1_output_region: ClassVar[str] = "blackwell.swap_ab_fc12.fc1_output" + fc1_output_sf_region: ClassVar[str] = "blackwell.swap_ab_fc12.fc1_output_sf" + fc1_done_counter_region: ClassVar[str] = "blackwell.swap_ab_fc12.fc1_done_counter" + epilogue_warp_ids: ClassVar[Tuple[int, int, int, int]] = (0, 1, 2, 3) + mma_warp_id: ClassVar[int] = 4 + tma_a_warp_id: ClassVar[int] = 5 + tma_b_warp_id: ClassVar[int] = 6 + scheduler_warp_id: ClassVar[int] = 7 + threads_per_cta: ClassVar[int] = 8 * 32 + + @classmethod + def problem_desc_require(cls) -> dict[str, object]: + return { + "expert_count": StaticOrRuntimeIntegerType, + "intermediate_gateup_size": StaticOrRuntimeIntegerType, + "hidden_size": StaticOrRuntimeIntegerType, + "quant_kind": str, + "a_major_mode": OperandMajorMode, + "b_major_mode": OperandMajorMode, + "combine_format": CombineFormat, + "gate_up_clamp": Optional[float], + "situ_beta": OptionalRequirement(Optional[float]), + "situ_linear_beta": OptionalRequirement(Optional[float]), + } + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return { + "mma_tiler_mnk": tuple, + "cluster_shape_mn": tuple, + "use_2cta_instrs": bool, + "group_hint": int, + "token_padding_block": int, + "sf_padding_block": int, + "work_id_mode": str, + "max_tokens": int, + "launch_cluster_count": int, + "fc2_use_bulk": bool, + "fc1_epi_flag_batch": int, + "fc2_epi_flag_batch": int, + } + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + self.expert_count = problem_desc["expert_count"] + self.intermediate_gateup_size = problem_desc["intermediate_gateup_size"] + self.hidden_size = problem_desc["hidden_size"] + self.quant_kind = QuantKind(problem_desc["quant_kind"]) + self.a_dtype = self.quant_kind.weight_dtype + self.b_dtype = self.quant_kind.activation_dtype + self.sf_dtype = self.quant_kind.sf_dtype + self.sf_vec_size = self.quant_kind.sf_vec_size + self.acc_dtype = tcgen05_block_scaled_acc_dtype + self.a_major_mode = problem_desc["a_major_mode"] + self.b_major_mode = problem_desc["b_major_mode"] + self.combine_format = problem_desc["combine_format"] + self.gate_up_clamp = problem_desc["gate_up_clamp"] + self.situ_beta = problem_desc.get("situ_beta") + self.situ_linear_beta = problem_desc.get("situ_linear_beta") + + self.mma_tiler_mnk = impl_desc["mma_tiler_mnk"] + self.cluster_shape_mn = impl_desc["cluster_shape_mn"] + self.use_2cta_instrs = impl_desc["use_2cta_instrs"] + self.group_hint = impl_desc["group_hint"] + self.token_padding_block = impl_desc["token_padding_block"] + self.sf_padding_block = impl_desc["sf_padding_block"] + self.work_id_mode: WorkIdAcquisitionMode = impl_desc["work_id_mode"] + self.max_tokens = impl_desc["max_tokens"] + self.launch_cluster_count = impl_desc["launch_cluster_count"] + self.fc2_use_bulk = impl_desc["fc2_use_bulk"] + self.fc1_epi_flag_batch = impl_desc["fc1_epi_flag_batch"] + self.fc2_epi_flag_batch = impl_desc["fc2_epi_flag_batch"] + + self.occupancy = 1 + self.architecture = "sm_100" + + self._validate_geometry() + mma_cta_count = 2 if self.use_2cta_instrs else 1 + mma_instruction = Tcgen05MmaInstruction( + a_type=self.a_dtype, + b_type=self.b_dtype, + acc_type=self.acc_dtype, + instruction_mnk=(self.mma_tiler_mnk[0], self.mma_tiler_mnk[1], self.quant_kind.instruction_k("1x")), + participates=mma_cta_count, + sfa_type=self.sf_dtype, + sfb_type=self.sf_dtype, + sf_vec_size=self.sf_vec_size, + ) + tmem_plan = make_tcgen05_tmem_plan(mma_instruction, self.architecture, self.mma_tiler_mnk) + resolved_impl_desc = ImplDesc( + { + **impl_desc, + "hint": self.group_hint, + "tmem_plan": tmem_plan, + "is_swap_ab": True, + "num_scheduler_consumer_threads": self.threads_per_cta - 32, + "num_accumulator_consumer_warps_per_cta": len(self.epilogue_warp_ids), + "communication_enabled": False, + } + ) + self.scheduler = BlackwellFusedFc12Scheduler(problem_desc, resolved_impl_desc) + self.epilogue = SwapABGatedActEpilogue(problem_desc, resolved_impl_desc) + if self.epilogue.fc1_output_dtype is not self.b_dtype: + raise ValueError("Epilogue FC1 output dtype must match the Mainloop B dtype.") + if self.epilogue.fc1_output_sf_dtype is not self.sf_dtype: + raise ValueError("Epilogue FC1 output scale dtype must match the Mainloop scale dtype.") + if self.epilogue.sf_vec_size != self.sf_vec_size: + raise ValueError("Epilogue and Mainloop scale vector sizes must match.") + self._device_workspace = self._build_device_workspace() + self._mainloop, self._smem_workspace = self._build_mainloop_and_smem(problem_desc, resolved_impl_desc) + + def _validate_geometry(self) -> None: + static_expert_dimensions = ( + isinstance(self.expert_count, int), + isinstance(self.intermediate_gateup_size, int), + isinstance(self.hidden_size, int), + ) + if any(static_expert_dimensions) and not all(static_expert_dimensions): + raise ValueError("FC12 expert dimensions must be either all static or all runtime.") + if not all(static_expert_dimensions): + raise NotImplementedError("The Lean FC12 Kernel currently requires static expert dimensions.") + if self.expert_count <= 0 or self.intermediate_gateup_size <= 0 or self.hidden_size <= 0: + raise ValueError("FC12 expert dimensions must be positive.") + if self.launch_cluster_count <= 0: + raise ValueError("launch_cluster_count must be positive.") + if self.intermediate_gateup_size % 2 != 0: + raise ValueError("The SwiGLU intermediate dimension must be even.") + if self.max_tokens <= 0: + raise ValueError("max_tokens must be positive.") + if self.token_padding_block <= 0: + raise ValueError("token_padding_block must be positive.") + if self.token_padding_block % 64 != 0: + raise ValueError("token_padding_block must be a multiple of 64.") + if self.sf_vec_size <= 0: + raise ValueError("sf_vec_size must be positive.") + if self.quant_kind.needs_unpack_tma(self.architecture): + for name, extent in ( + ("hidden_size", self.hidden_size), + ("intermediate_downproj", self.intermediate_gateup_size // 2), + ): + if extent % 128 != 0: + raise ValueError( + f"{self.quant_kind} loads its fp4 weight through the unpacking TMA, which " + f"requires {name} to be a multiple of 128, got {extent}." + ) + mma_m, mma_n, mma_k = self.mma_tiler_mnk + cluster_m, cluster_n = self.cluster_shape_mn + expected_mma_m = 256 if self.use_2cta_instrs else 128 + if mma_m != expected_mma_m: + raise ValueError(f"mma_tiler M must be {expected_mma_m}, got {mma_m}.") + if mma_n not in (64, 128, 256): + raise ValueError(f"mma_tiler N must be 64, 128, or 256, got {mma_n}.") + if mma_k % (self.sf_vec_size * 4) != 0: + raise ValueError("mma_tiler K must be divisible by four scale-factor vectors.") + if self.intermediate_gateup_size % (self.sf_vec_size * 4) != 0: + raise ValueError("The intermediate dimension must be divisible by four scale-factor vectors.") + if cluster_n != 1: + raise ValueError(f"The swap-AB FC12 path requires cluster N=1, got {cluster_n}.") + if cluster_m <= 0 or cluster_m > 16 or cluster_m & (cluster_m - 1): + raise ValueError("cluster M must be a power of two no greater than 16.") + if self.use_2cta_instrs and cluster_m % 2 != 0: + raise ValueError("Two-CTA MMA requires an even cluster M.") + + def _build_device_workspace(self) -> DeviceWorkspace: + intermediate_downproj = self.intermediate_gateup_size // 2 + sf_column_count = round_up(ceil_div(intermediate_downproj, self.sf_vec_size), 4) + max_sf_rows = self.max_tokens + self.expert_count * self.sf_padding_block + counter_slot_count = ceil_div(self.max_tokens, self.mma_tiler_mnk[1]) + self.expert_count + + device_workspace = DeviceWorkspace() + device_workspace.register( + self.fc1_output_region, + self.epilogue.fc1_output_dtype, + (self.max_tokens, intermediate_downproj), + buffer_space="local", + mem_order=(1, 0), + byte_alignment=128, + ) + device_workspace.register( + self.fc1_output_sf_region, + self.epilogue.fc1_output_sf_dtype, + (max_sf_rows, sf_column_count), + buffer_space="local", + mem_order=(1, 0), + byte_alignment=128, + ) + device_workspace.register( + self.fc1_done_counter_region, + cutlass.Int32, + (counter_slot_count,), + buffer_space="local", + byte_alignment=16, + reset="tail_reset", + ) + self.scheduler.register_device_workspace(device_workspace) + device_workspace.finalize() + return device_workspace + + def get_workspace_sizes(self) -> Tuple[int, int]: + """Return required local and shared workspace bytes.""" + return self._device_workspace.local_and_shared_bytes + + @property + def require_zero_workspace_leading_bytes(self) -> Tuple[int, int]: + return self._device_workspace.require_zero_workspace_leading_bytes + + def name(self) -> str: + return "" + + def aot_compile(self): + return None + + def _build_mainloop_and_smem( + self, problem_desc: ProblemDesc, resolved_impl_desc: ImplDesc + ) -> Tuple[BlockScaledSwapAbFc12Mainloop, SmemWorkspace]: + smem_limit = utils.get_smem_capacity_in_bytes() // self.occupancy + smem_workspace = SmemWorkspace() + self.scheduler.register_smem_regions(smem_workspace) + self.epilogue.register_smem_regions(smem_workspace) + alignment_shift_bytes = max(region.byte_alignment for region in smem_workspace.regions()) + mainloop_smem_budget_bytes = smem_limit - smem_workspace.estimate_total_bytes() - alignment_shift_bytes + if mainloop_smem_budget_bytes <= 0: + raise ValueError("Fixed kernel components leave no SMEM budget for the Mainloop.") + + mainloop_impl_desc = ImplDesc({**resolved_impl_desc, "mainloop_smem_budget_bytes": mainloop_smem_budget_bytes}) + mainloop = BlockScaledSwapAbFc12Mainloop(problem_desc, mainloop_impl_desc) + + mainloop.register_smem_regions(smem_workspace) + smem_workspace.finalize(max_bytes=smem_limit) + return mainloop, smem_workspace + + @cute.jit + def __call__( + self, + activation: cute.Tensor, # (tokens, hidden) + fc1_weight: cute.Tensor, # (experts, hidden, intermediate_gateup) + activation_sf: cute.Tensor, # (padded_tokens, padded_hidden // sf_vec_size) + fc1_weight_sf: cute.Tensor, # (experts, packed_fc1_sf) + fc2_weight: cute.Tensor, # (experts, intermediate_downproj, hidden) + fc2_weight_sf: cute.Tensor, # (experts, packed_fc2_sf) + fc2_output: cute.Tensor, # (tokens, topk, hidden) + local_workspace: cute.Pointer, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + expert_token_sizes: Optional[cute.Tensor] = None, # (experts,) + expert_token_prefix_sum: Optional[cute.Tensor] = None, # (experts,) + topk_scores: Optional[cute.Tensor] = None, # (tokens,) + fc1_alpha: Optional[cute.Tensor] = None, + fc2_alpha: Optional[cute.Tensor] = None, + fc1_norm_const: Optional[cute.Tensor] = None, + ) -> None: + """Prepare and launch the lean fused FC1+FC2 kernel. + + The caller must zero the workspace prefix reported by + require_zero_workspace_leading_bytes before its first launch. + """ + if cutlass.const_expr((expert_token_sizes is None) == (expert_token_prefix_sum is None)): + raise ValueError("Exactly one of expert_token_sizes and expert_token_prefix_sum must be provided.") + # The quant kind fixes every operand format, so the tensors are where a caller mismatch can + # still show up. See BlockScaledSwapAbMegaMoeKernel.__call__. + for operand_name, operand, expected_dtype in ( + ("activation", activation, self.b_dtype), + ("activation_sf", activation_sf, self.sf_dtype), + ("fc1_weight", fc1_weight, self.a_dtype), + ("fc1_weight_sf", fc1_weight_sf, self.sf_dtype), + ("fc2_weight", fc2_weight, self.a_dtype), + ("fc2_weight_sf", fc2_weight_sf, self.sf_dtype), + ): + if cutlass.const_expr(operand.element_type is not expected_dtype): + raise TypeError( + f"{self.quant_kind} requires {operand_name} to be {expected_dtype}, got {operand.element_type}." + ) + if cutlass.const_expr(not self.quant_kind.uses_global_scale): + for scalar_name, scalar in ( + ("fc1_alpha", fc1_alpha), + ("fc2_alpha", fc2_alpha), + ("fc1_norm_const", fc1_norm_const), + ): + if cutlass.const_expr(scalar is not None): + raise ValueError( + f"{self.quant_kind} folds its rescale into the e8m0 scale factors, " + f"so {scalar_name} must be None." + ) + + def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: + return cute.make_tensor(tensor.iterator, cute.make_layout(shape, stride=tensor.stride)) + + self._device_workspace.assign_device_members(local_workspace) + fc1_output = self._device_workspace.tensor(self.fc1_output_region) + fc1_output_sf = self._device_workspace.tensor(self.fc1_output_sf_region) + + if cutlass.const_expr(isinstance(self.expert_count, int)): + experts = self.expert_count + else: + experts = fc1_weight.shape[0] + if cutlass.const_expr(isinstance(self.hidden_size, int)): + hidden = self.hidden_size + else: + hidden = fc1_weight.shape[1] + if cutlass.const_expr(isinstance(self.intermediate_gateup_size, int)): + intermediate_gateup = self.intermediate_gateup_size + else: + intermediate_gateup = fc1_weight.shape[2] + intermediate_downproj = intermediate_gateup // 2 + fc1_weight = rewrite_tensor_shape(fc1_weight, (experts, hidden, intermediate_gateup)) + fc2_weight = rewrite_tensor_shape(fc2_weight, (experts, intermediate_downproj, hidden)) + activation = rewrite_tensor_shape(activation, (activation.shape[0], hidden)) + if cutlass.const_expr(isinstance(activation.shape[0], int)): + if cutlass.const_expr(activation.shape[0] > self.max_tokens): + raise ValueError("activation rows exceed the configured max_tokens.") + fc2_output = rewrite_tensor_shape(fc2_output, (fc2_output.shape[0], fc2_output.shape[1], hidden)) + + singleton = cutlass.Int32(1) + token_rows = activation.shape[0] + + fc1_a = cute.make_tensor( + fc1_weight.iterator, + cute.make_layout( + (cutlass.Int32(intermediate_gateup), cutlass.Int32(hidden), cutlass.Int32(experts)), + stride=(fc1_weight.stride[2], fc1_weight.stride[1], fc1_weight.stride[0]), + ), + ) + fc1_b = cute.make_tensor( + activation.iterator, + cute.make_layout( + (token_rows, cutlass.Int32(hidden), singleton), stride=(activation.stride[0], activation.stride[1], 0) + ), + ) + fc1_output_gemm = cute.make_tensor( + fc1_output.iterator, + cute.make_layout( + (token_rows, cutlass.Int32(intermediate_downproj), singleton), + stride=(fc1_output.stride[0], fc1_output.stride[1], 0), + ), + ) + + sf_vec_size = self.sf_vec_size + padded_token_rows = activation_sf.shape[0] + padded_hidden = activation_sf.shape[1] * sf_vec_size + fc1_sfb = cute.make_tensor( + activation_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (padded_token_rows, cutlass.Int32(padded_hidden), singleton), sf_vec_size + ), + ) + padded_intermediate_gateup = cute.round_up(intermediate_gateup, sf_vec_size * 4) + expected_fc1_weight_sf_columns = padded_intermediate_gateup * padded_hidden // sf_vec_size + if cutlass.const_expr( + isinstance(fc1_weight_sf.shape[1], int) and isinstance(expected_fc1_weight_sf_columns, int) + ): + if cutlass.const_expr(fc1_weight_sf.shape[1] != expected_fc1_weight_sf_columns): + raise ValueError("fc1_weight_sf has an incompatible column count.") + fc1_sfa = cute.make_tensor( + fc1_weight_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (cutlass.Int32(padded_intermediate_gateup), cutlass.Int32(padded_hidden), cutlass.Int32(experts)), + sf_vec_size, + ), + ) + + experts_fc2, intermediate_fc2, hidden_fc2 = fc2_weight.shape + fc2_a = cute.make_tensor( + fc2_weight.iterator, + cute.make_layout( + (cutlass.Int32(hidden_fc2), cutlass.Int32(intermediate_fc2), cutlass.Int32(experts_fc2)), + stride=(fc2_weight.stride[2], fc2_weight.stride[1], fc2_weight.stride[0]), + ), + ) + fc2_b = fc1_output_gemm + padded_fc2_hidden = cute.round_up(hidden_fc2, 128) + padded_fc2_intermediate = cute.round_up(intermediate_fc2, sf_vec_size * 4) + expected_fc2_weight_sf_columns = padded_fc2_hidden * padded_fc2_intermediate // sf_vec_size + if cutlass.const_expr( + isinstance(fc2_weight_sf.shape[1], int) and isinstance(expected_fc2_weight_sf_columns, int) + ): + if cutlass.const_expr(fc2_weight_sf.shape[1] != expected_fc2_weight_sf_columns): + raise ValueError("fc2_weight_sf has an incompatible column count.") + fc2_sfa = cute.make_tensor( + fc2_weight_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (cutlass.Int32(padded_fc2_hidden), cutlass.Int32(padded_fc2_intermediate), cutlass.Int32(experts_fc2)), + sf_vec_size, + ), + ) + fc2_sfb = cute.make_tensor( + fc1_output_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (fc1_output_sf.shape[0], cutlass.Int32(fc1_output_sf.shape[1] * sf_vec_size), singleton), sf_vec_size + ), + ) + + self._mainloop.materialize_codegen_members() + ( + fc1_tma_a_tensor, + fc1_tma_a_atom, + fc1_tma_sfa_tensor, + fc1_tma_sfa_atom, + fc2_tma_a_tensor, + fc2_tma_a_atom, + fc2_tma_sfa_tensor, + fc2_tma_sfa_atom, + fc1_tma_b_tensor, + fc1_tma_b_atom_or_pair, + fc1_tma_sfb_tensor, + fc1_tma_sfb_atom_or_pair, + fc2_tma_b_tensor, + fc2_tma_b_atom_or_pair, + fc2_tma_sfb_tensor, + fc2_tma_sfb_atom_or_pair, + ) = self._mainloop.prepare_tma_load_params( + fc1_a=fc1_a, + fc1_b=fc1_b, + fc1_sfa=fc1_sfa, + fc1_sfb=fc1_sfb, + fc2_a=fc2_a, + fc2_b=fc2_b, + fc2_sfa=fc2_sfa, + fc2_sfb=fc2_sfb, + ) + + (fc1_output_tma_atom, fc1_output_tma_tensor, fc2_output_tma_atom, fc2_output_tma_tensor) = ( + self.epilogue.prepare_tma_store_params(fc1_output_gemm, fc2_output) + ) + + grid = self.scheduler.get_grid_shape(max_active_clusters=max_active_clusters) + self._kernel( + fc1_tma_a_tensor, + fc1_tma_a_atom, + fc1_tma_sfa_tensor, + fc1_tma_sfa_atom, + fc2_tma_a_tensor, + fc2_tma_a_atom, + fc2_tma_sfa_tensor, + fc2_tma_sfa_atom, + fc1_tma_b_tensor, + fc1_tma_b_atom_or_pair, + fc1_tma_sfb_tensor, + fc1_tma_sfb_atom_or_pair, + fc2_tma_b_tensor, + fc2_tma_b_atom_or_pair, + fc2_tma_sfb_tensor, + fc2_tma_sfb_atom_or_pair, + fc1_output_tma_atom, + fc1_output_tma_tensor, + fc2_output_tma_atom, + fc2_output_tma_tensor, + fc2_output, + expert_token_sizes, + expert_token_prefix_sum, + (experts, intermediate_gateup, hidden), + local_workspace, + fc1_alpha, + fc2_alpha, + fc1_norm_const, + topk_scores, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + min_blocks_per_mp=self.occupancy, + ) + + @cute.kernel + def _kernel( + self, + fc1_tma_a_tensor: cute.Tensor, + fc1_tma_a_atom: cute.CopyAtom, + fc1_tma_sfa_tensor: cute.Tensor, + fc1_tma_sfa_atom: cute.CopyAtom, + fc2_tma_a_tensor: cute.Tensor, + fc2_tma_a_atom: cute.CopyAtom, + fc2_tma_sfa_tensor: cute.Tensor, + fc2_tma_sfa_atom: cute.CopyAtom, + fc1_tma_b_tensor: cute.Tensor, + fc1_tma_b_atom_or_pair: TmaAtomOrPair, + fc1_tma_sfb_tensor: cute.Tensor, + fc1_tma_sfb_atom_or_pair: TmaAtomOrPair, + fc2_tma_b_tensor: cute.Tensor, + fc2_tma_b_atom_or_pair: TmaAtomOrPair, + fc2_tma_sfb_tensor: cute.Tensor, + fc2_tma_sfb_atom_or_pair: TmaAtomOrPair, + fc1_output_tma_atom: cute.CopyAtom, + fc1_output_tma_tensor: cute.Tensor, + fc2_output_tma_atom: Optional[cute.CopyAtom], + fc2_output_tma_tensor: Optional[cute.Tensor], + fc2_output: cute.Tensor, + expert_token_sizes: Optional[cute.Tensor], + expert_token_prefix_sum: Optional[cute.Tensor], + actual_expert_shape: Tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32], + local_workspace: cute.Pointer, + fc1_alpha: Optional[cute.Tensor], + fc2_alpha: Optional[cute.Tensor], + fc1_norm_const: Optional[cute.Tensor], + topk_scores: Optional[cute.Tensor], + ): + """Compose Scheduler, Mainloop, and Epilogue.""" + self._mainloop.materialize_codegen_members() + storage_type = self._smem_workspace.storage_class() + smem_allocator = utils.SmemAllocator() + storage = smem_allocator.allocate(storage_type) + smem_base = storage.buffer.data_ptr() + self._device_workspace.assign_device_members(local_workspace) + fc1_done_counter = self._device_workspace.tensor(self.fc1_done_counter_region) + fc1_output_sf_storage = self._device_workspace.tensor(self.fc1_output_sf_region) + if cutlass.const_expr(isinstance(self.hidden_size, int)): + hidden = self.hidden_size + else: + hidden = actual_expert_shape[2] + if cutlass.const_expr(isinstance(self.intermediate_gateup_size, int)): + intermediate_gateup = self.intermediate_gateup_size + else: + intermediate_gateup = actual_expert_shape[1] + fc1_output_sf = cute.make_tensor( + fc1_output_sf_storage.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (fc1_output_sf_storage.shape[0], intermediate_gateup // 2, cutlass.Int32(1)), self.epilogue.sf_vec_size + ), + ) + + thread_idx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.make_warp_uniform(thread_idx // 32) + block_idx = cute.arch.block_idx() + + cta_rank_in_cluster = cute.arch.block_idx_in_cluster() + cta_coord_in_cluster = ( + cta_rank_in_cluster % self.cluster_shape_mn[0], + cta_rank_in_cluster // self.cluster_shape_mn[0], + cutlass.Int32(0), + ) + + # Standalone FC12 binds the preferred layout before pipeline creation. + self._mainloop.assign_device_members( + self._smem_workspace, smem_base, cta_coord_in_cluster, Boolean(False), hidden, intermediate_gateup + ) + ab_pipeline = self._mainloop.create_ab_pipeline(self._smem_workspace, smem_base) + tma_a_pipeline_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self._mainloop.num_ab_pipeline_stages + ) + tma_b_pipeline_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self._mainloop.num_ab_pipeline_stages + ) + mma_pipeline_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self._mainloop.num_ab_pipeline_stages + ) + acc_pipeline = self._mainloop.create_acc_pipeline(self._smem_workspace, smem_base) + tmem_allocator = self._mainloop.create_tmem_allocator( + self._smem_workspace, smem_base, allocator_warp_id=self.epilogue_warp_ids[0] + ) + self.scheduler.assign_device_members( + expert_token_sizes=expert_token_sizes, + expert_token_prefix_sum=expert_token_prefix_sum, + actual_expert_shape=actual_expert_shape, + block_idx=block_idx, + smem_workspace=self._smem_workspace, + smem_base=smem_base, + device_workspace=self._device_workspace, + ) + scheduler = self.scheduler + + fc2_spin_threshold = (intermediate_gateup + self._mainloop.cta_tile_m - 1) // self._mainloop.cta_tile_m + kernel_extension = BlockScaledSwapAbFc12Extension( + sf_vec_size=self.sf_vec_size, + fc1_done_counter_pointer=fc1_done_counter.iterator, + fc2_spin_threshold=fc2_spin_threshold, + fc1_ready_counter_pointer=None, + ) + optional_epilogue_args = GatedActEpilogueArgs( + fc1_alpha=fc1_alpha, fc2_alpha=fc2_alpha, fc1_norm_const=fc1_norm_const, topk_scores=topk_scores + ) + + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + if warp_idx == self.scheduler_warp_id: + work_tile = scheduler.gen_next_work() + while work_tile.is_valid_tile: + scheduler.publish_work(kernel_extension.prepare_work_tile(work_tile)) + work_tile = scheduler.gen_next_work() + scheduler.publish_work(work_tile) + scheduler.produce_tail() + + if warp_idx == self.tma_a_warp_id: + sched_consumer = scheduler.make_consumer() + self._mainloop.run_tma_a( + fc1_tma_a_tensor=fc1_tma_a_tensor, + fc1_tma_a_atom=fc1_tma_a_atom, + fc1_tma_sfa_tensor=fc1_tma_sfa_tensor, + fc1_tma_sfa_atom=fc1_tma_sfa_atom, + fc2_tma_a_tensor=fc2_tma_a_tensor, + fc2_tma_a_atom=fc2_tma_a_atom, + fc2_tma_sfa_tensor=fc2_tma_sfa_tensor, + fc2_tma_sfa_atom=fc2_tma_sfa_atom, + ab_pipeline=ab_pipeline, + ab_pipeline_state=tma_a_pipeline_state, + sched_consumer=sched_consumer, + kernel_extension=kernel_extension, + ) + + if warp_idx == self.tma_b_warp_id: + sched_consumer = scheduler.make_consumer() + self._mainloop.run_tma_b( + fc1_tma_b_tensor=fc1_tma_b_tensor, + fc1_tma_b_atom_or_pair=fc1_tma_b_atom_or_pair, + fc1_tma_sfb_tensor=fc1_tma_sfb_tensor, + fc1_tma_sfb_atom_or_pair=fc1_tma_sfb_atom_or_pair, + fc2_tma_b_tensor=fc2_tma_b_tensor, + fc2_tma_b_atom_or_pair=fc2_tma_b_atom_or_pair, + fc2_tma_sfb_tensor=fc2_tma_sfb_tensor, + fc2_tma_sfb_atom_or_pair=fc2_tma_sfb_atom_or_pair, + ab_pipeline=ab_pipeline, + ab_pipeline_state=tma_b_pipeline_state, + sched_consumer=sched_consumer, + kernel_extension=kernel_extension, + fc1_done_counter_pointer=fc1_done_counter.iterator, + fc2_spin_threshold=fc2_spin_threshold, + ) + + if warp_idx == self.mma_warp_id: + sched_consumer = scheduler.make_consumer() + self._mainloop.run_mma( + tmem_allocator=tmem_allocator, + ab_pipeline=ab_pipeline, + ab_pipeline_state=mma_pipeline_state, + acc_pipeline=acc_pipeline, + sched_consumer=sched_consumer, + ) + + if warp_idx < len(self.epilogue_warp_ids): + sched_consumer = scheduler.make_consumer() + tmem_allocator.allocate(self._mainloop.num_tmem_alloc_cols) + tmem_allocator.wait_for_alloc() + tmem_pointer = tmem_allocator.retrieve_ptr(self.acc_dtype) + self.epilogue.run( + self._smem_workspace, + smem_base, + tmem_pointer, + acc_pipeline, + sched_consumer, + kernel_extension, + fc1_output_tma_atom, + fc1_output_tma_tensor, + fc1_output_sf, + fc2_output_tma_atom, + fc2_output_tma_tensor, + fc2_output, + fc1_done_counter, + thread_idx, + None, + None, + None, + None, + optional_epilogue_args, + ) + tmem_allocator.relinquish_alloc_permit() + tmem_allocator.free(tmem_allocator.retrieve_ptr(self.acc_dtype), self._mainloop.num_tmem_alloc_cols) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_mainloop.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_mainloop.py index c04f39fc3686..93d61ff8f69c 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_mainloop.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_mainloop.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Mainloop component for the block-scaled swap-AB FC12 kernel.""" from typing import ClassVar, Optional, Tuple @@ -14,13 +11,7 @@ from cutlass.cute.nvgpu import OperandMajorMode, cpasync, tcgen05 from cutlass.cutlass_dsl import Boolean, Int32 -from .....api import ( - ImplDesc, - KernelComponent, - OptionalRequirement, - ProblemDesc, - StaticOrRuntimeIntegerType, -) +from .....api import ImplDesc, KernelComponent, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType from .....helpers.cute_py_helpers import ( Tcgen05MmaInstruction, make_smem_layouts, @@ -45,9 +36,7 @@ @cute.jit -def _make_executable_tma_atom_for_cluster( - atom_or_pair: TmaAtomOrPair, is_fallback_cluster: Boolean -) -> cute.CopyAtom: +def _make_executable_tma_atom_for_cluster(atom_or_pair: TmaAtomOrPair, is_fallback_cluster: Boolean) -> cute.CopyAtom: """CuTeDSL < 4.7 WAR: materialize exec atoms inside the selecting scf.if branches.""" if cutlass.const_expr(isinstance(atom_or_pair, tuple)): executable_atom = make_executable_tma_atom(atom_or_pair[0]) @@ -62,22 +51,14 @@ def _make_executable_tma_atom_for_cluster( class BlockScaledSwapAbFc12Mainloop(KernelComponent): """Own all load, MMA, pipeline, SMEM, and TMEM mainloop state.""" - ab_pipeline_mbarriers_region: ClassVar[str] = ( - "blackwell.swap_ab_fc12.mainloop.ab_pipeline_mbarriers" - ) + ab_pipeline_mbarriers_region: ClassVar[str] = "blackwell.swap_ab_fc12.mainloop.ab_pipeline_mbarriers" a_smem_tensor_region: ClassVar[str] = "blackwell.swap_ab_fc12.mainloop.a_smem_tensor" b_smem_tensor_region: ClassVar[str] = "blackwell.swap_ab_fc12.mainloop.b_smem_tensor" sfa_smem_tensor_region: ClassVar[str] = "blackwell.swap_ab_fc12.mainloop.sfa_smem_tensor" sfb_smem_tensor_region: ClassVar[str] = "blackwell.swap_ab_fc12.mainloop.sfb_smem_tensor" - acc_pipeline_mbarriers_region: ClassVar[str] = ( - "blackwell.swap_ab_fc12.mainloop.acc_pipeline_mbarriers" - ) - tmem_holding_buffer_region: ClassVar[str] = ( - "blackwell.swap_ab_fc12.mainloop.tmem_holding_buffer" - ) - tmem_deallocation_mbarrier_region: ClassVar[str] = ( - "blackwell.swap_ab_fc12.mainloop.tmem_deallocation_mbarrier" - ) + acc_pipeline_mbarriers_region: ClassVar[str] = "blackwell.swap_ab_fc12.mainloop.acc_pipeline_mbarriers" + tmem_holding_buffer_region: ClassVar[str] = "blackwell.swap_ab_fc12.mainloop.tmem_holding_buffer" + tmem_deallocation_mbarrier_region: ClassVar[str] = "blackwell.swap_ab_fc12.mainloop.tmem_deallocation_mbarrier" tmem_allocation_barrier_id: ClassVar[int] = 2 @classmethod @@ -120,9 +101,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.fallback_cluster_shape_mn = impl_desc.get("fallback_cluster_shape_mn") self.use_2cta_instrs = impl_desc["use_2cta_instrs"] self.mainloop_smem_budget_bytes = impl_desc["mainloop_smem_budget_bytes"] - self.num_accumulator_consumer_warps_per_cta = impl_desc[ - "num_accumulator_consumer_warps_per_cta" - ] + self.num_accumulator_consumer_warps_per_cta = impl_desc["num_accumulator_consumer_warps_per_cta"] self._validate_configuration() self.architecture = "sm_100" @@ -141,11 +120,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: tmem_plan.accumulator_stage_columns - tmem_plan.accumulator_stage_stride_columns ) self.num_accumulator_tmem_cols = tmem_plan.accumulator_columns - self.accumulator_shape = ( - self.cta_tile_m, - self.cta_tile_n, - tmem_plan.accumulator_stage_count, - ) + self.accumulator_shape = (self.cta_tile_m, self.cta_tile_n, tmem_plan.accumulator_stage_count) self.accumulator_stride = (1 << 16, 1, tmem_plan.accumulator_stage_stride_columns) self.overlapping_accum = self.accumulator_overlap_columns > 0 @@ -153,11 +128,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: a_type=self.a_dtype, b_type=self.b_dtype, acc_type=self.acc_dtype, - instruction_mnk=( - self.mma_tiler_mnk[0], - self.mma_tiler_mnk[1], - self.quant_kind.instruction_k("1x"), - ), + instruction_mnk=(self.mma_tiler_mnk[0], self.mma_tiler_mnk[1], self.quant_kind.instruction_k("1x")), participates=self.mma_cta_count, sfa_type=self.sf_dtype, sfb_type=self.sf_dtype, @@ -165,12 +136,8 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: ) # A mixed-width pair reaches SMEM through the unpacking TMA, so the narrow operand's SMEM # image is byte-per-element even though the MMA still sees its logical type. - self.a_smem_alloc_dtype = tcgen05_smem_alloc_type( - self.a_dtype, self.b_dtype, self.architecture - ) - self.b_smem_alloc_dtype = tcgen05_smem_alloc_type( - self.b_dtype, self.a_dtype, self.architecture - ) + self.a_smem_alloc_dtype = tcgen05_smem_alloc_type(self.a_dtype, self.b_dtype, self.architecture) + self.b_smem_alloc_dtype = tcgen05_smem_alloc_type(self.b_dtype, self.a_dtype, self.architecture) self.mma_tiler_sfb = (self.cta_tile_m, round_up(self.cta_tile_n, 128), self.mma_tile_k) self.cluster_layout_shape_vmnk = ( (self.mma_cta_count,), @@ -180,9 +147,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: ) self.cluster_layout_sfb_shape_vmnk = ((1,), *self.cluster_shape_mn, 1) self.resolved_fallback_cluster_shape_mn = ( - self.cluster_shape_mn - if self.fallback_cluster_shape_mn is None - else self.fallback_cluster_shape_mn + self.cluster_shape_mn if self.fallback_cluster_shape_mn is None else self.fallback_cluster_shape_mn ) self.fallback_cluster_layout_shape_vmnk = ( (self.mma_cta_count,), @@ -190,11 +155,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.resolved_fallback_cluster_shape_mn[1], 1, ) - self.fallback_cluster_layout_sfb_shape_vmnk = ( - (1,), - *self.resolved_fallback_cluster_shape_mn, - 1, - ) + self.fallback_cluster_layout_sfb_shape_vmnk = ((1,), *self.resolved_fallback_cluster_shape_mn, 1) self.is_mixed_cga = self.resolved_fallback_cluster_shape_mn != self.cluster_shape_mn self.num_mcast_ctas_a = self.cluster_layout_shape_vmnk[2] self.num_mcast_ctas_b = self.cluster_layout_shape_vmnk[1] @@ -214,9 +175,7 @@ def _validate_configuration(self) -> None: if any(dimension <= 0 for dimension in self.cluster_shape_mn): raise ValueError("cluster_shape_mn dimensions must be positive.") if self.mma_tiler_mnk[1] not in (64, 128, 256): - raise ValueError( - f"The component supports mma_tiler N in (64, 128, 256); got {self.mma_tiler_mnk[1]}." - ) + raise ValueError(f"The component supports mma_tiler N in (64, 128, 256); got {self.mma_tiler_mnk[1]}.") mma_cta_count = 2 if self.use_2cta_instrs else 1 if self.mma_tiler_mnk[0] % mma_cta_count != 0: raise ValueError("mma_tiler M must be divisible by the MMA CTA count.") @@ -240,9 +199,7 @@ def _validate_configuration(self) -> None: self.cluster_shape_mn, self.fallback_cluster_shape_mn ) ): - raise ValueError( - "Preferred cluster dimensions must be divisible by fallback dimensions." - ) + raise ValueError("Preferred cluster dimensions must be divisible by fallback dimensions.") if self.mainloop_smem_budget_bytes <= 0: raise ValueError("mainloop_smem_budget_bytes must be positive.") if self.num_accumulator_consumer_warps_per_cta <= 0: @@ -257,9 +214,7 @@ def _validate_configuration(self) -> None: def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: """Register the selected mainloop SMEM plan.""" - smem_workspace.register_mbarrier( - self.ab_pipeline_mbarriers_region, self.num_ab_pipeline_stages * 2 - ) + smem_workspace.register_mbarrier(self.ab_pipeline_mbarriers_region, self.num_ab_pipeline_stages * 2) smem_workspace.register_tensor( self.a_smem_tensor_region, self.a_smem_region.dtype, @@ -292,21 +247,13 @@ def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: swizzle=self.sfb_smem_region.swizzle, byte_alignment=self.sfb_smem_region.byte_alignment, ) - smem_workspace.register_mbarrier( - self.acc_pipeline_mbarriers_region, self.num_accumulator_pipeline_stages * 2 - ) + smem_workspace.register_mbarrier(self.acc_pipeline_mbarriers_region, self.num_accumulator_pipeline_stages * 2) smem_workspace.register_mbarrier(self.tmem_deallocation_mbarrier_region, 1) - smem_workspace.register_tensor( - self.tmem_holding_buffer_region, cutlass.Int32, (1,), byte_alignment=8 - ) + smem_workspace.register_tensor(self.tmem_holding_buffer_region, cutlass.Int32, (1,), byte_alignment=8) def _select_ab_stage_plan(self) -> None: stage_one_regions = make_smem_layouts( - self.mma_instruction, - self.mma_tiler_mnk, - 1, - (self.a_major_mode, self.b_major_mode), - self.architecture, + self.mma_instruction, self.mma_tiler_mnk, 1, (self.a_major_mode, self.b_major_mode), self.architecture ) a_region, b_region, sfa_region, sfb_region = stage_one_regions @@ -323,32 +270,23 @@ def packed_bytes(region, dtype) -> int: mbarrier_bytes = int(cutlass.Int64.width) // 8 ab_stage_cost_bytes = self.ab_stage_payload_bytes + 2 * mbarrier_bytes plan_tail_bytes = (2 * self.num_accumulator_pipeline_stages + 2) * mbarrier_bytes - self.num_ab_pipeline_stages = ( - self.mainloop_smem_budget_bytes - plan_tail_bytes - ) // ab_stage_cost_bytes + self.num_ab_pipeline_stages = (self.mainloop_smem_budget_bytes - plan_tail_bytes) // ab_stage_cost_bytes if self.num_ab_pipeline_stages < 1: raise ValueError( "One AB pipeline stage needs " f"{ab_stage_cost_bytes + plan_tail_bytes} bytes, exceeding the " f"{self.mainloop_smem_budget_bytes}-byte mainloop budget." ) - (self.a_smem_region, self.b_smem_region, self.sfa_smem_region, self.sfb_smem_region) = ( - make_smem_layouts( - self.mma_instruction, - self.mma_tiler_mnk, - self.num_ab_pipeline_stages, - (self.a_major_mode, self.b_major_mode), - self.architecture, - ) + (self.a_smem_region, self.b_smem_region, self.sfa_smem_region, self.sfb_smem_region) = make_smem_layouts( + self.mma_instruction, + self.mma_tiler_mnk, + self.num_ab_pipeline_stages, + (self.a_major_mode, self.b_major_mode), + self.architecture, ) self.selected_smem_bytes = sum( region.nbytes - for region in ( - self.a_smem_region, - self.b_smem_region, - self.sfa_smem_region, - self.sfb_smem_region, - ) + for region in (self.a_smem_region, self.b_smem_region, self.sfa_smem_region, self.sfb_smem_region) ) def materialize_codegen_members(self) -> None: @@ -376,12 +314,10 @@ def materialize_codegen_members(self) -> None: cute.make_layout((*self.cluster_shape_mn, 1)), (self.tiled_mma_sfb.thr_id.shape,) ) self.fallback_cluster_layout_vmnk = cute.tiled_divide( - cute.make_layout((*self.resolved_fallback_cluster_shape_mn, 1)), - (self.tiled_mma.thr_id.shape,), + cute.make_layout((*self.resolved_fallback_cluster_shape_mn, 1)), (self.tiled_mma.thr_id.shape,) ) self.fallback_cluster_layout_sfb_vmnk = cute.tiled_divide( - cute.make_layout((*self.resolved_fallback_cluster_shape_mn, 1)), - (self.tiled_mma_sfb.thr_id.shape,), + cute.make_layout((*self.resolved_fallback_cluster_shape_mn, 1)), (self.tiled_mma_sfb.thr_id.shape,) ) if self.cluster_layout_vmnk.shape != self.cluster_layout_shape_vmnk: raise ValueError( @@ -398,10 +334,7 @@ def materialize_codegen_members(self) -> None: f"Fallback cluster layout mismatch: Python {self.fallback_cluster_layout_shape_vmnk}, " f"CuTe {self.fallback_cluster_layout_vmnk.shape}." ) - if ( - self.fallback_cluster_layout_sfb_vmnk.shape - != self.fallback_cluster_layout_sfb_shape_vmnk - ): + if self.fallback_cluster_layout_sfb_vmnk.shape != self.fallback_cluster_layout_sfb_shape_vmnk: raise ValueError( f"Fallback SFB cluster layout mismatch: Python {self.fallback_cluster_layout_sfb_shape_vmnk}, " f"CuTe {self.fallback_cluster_layout_sfb_vmnk.shape}." @@ -429,11 +362,8 @@ def materialize_codegen_members(self) -> None: ("SFA", self.sfa_smem_layout_staged, self.sfa_smem_region), ("SFB", self.sfb_smem_layout_staged, self.sfb_smem_region), ): - if ( - materialized_layout.shape != planned_region.shape - or not strides_equal_ignoring_singletons( - planned_region.shape, materialized_layout.stride, planned_region.stride - ) + if materialized_layout.shape != planned_region.shape or not strides_equal_ignoring_singletons( + planned_region.shape, materialized_layout.stride, planned_region.stride ): raise ValueError( f"{tensor_name} SMEM plan mismatch: Python " @@ -459,27 +389,15 @@ def prepare_tma_load_params( sfa_stage_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0)) sfb_stage_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0)) - a_operation = sm100_utils.cluster_shape_to_tma_atom_A( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - b_operation = sm100_utils.cluster_shape_to_tma_atom_B( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - sfa_operation = sm100_utils.cluster_shape_to_tma_atom_A( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - sfb_operation = sm100_utils.cluster_shape_to_tma_atom_SFB( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) + a_operation = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, self.tiled_mma.thr_id) + b_operation = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, self.tiled_mma.thr_id) + sfa_operation = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, self.tiled_mma.thr_id) + sfb_operation = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, self.tiled_mma.thr_id) # An operand whose SMEM container is wider than its element type must be loaded by the # unpacking TMA variant; naming the container as internal_type is what selects it. - a_internal_type = ( - self.a_smem_alloc_dtype if self.a_smem_alloc_dtype is not self.a_dtype else None - ) - b_internal_type = ( - self.b_smem_alloc_dtype if self.b_smem_alloc_dtype is not self.b_dtype else None - ) + a_internal_type = self.a_smem_alloc_dtype if self.a_smem_alloc_dtype is not self.a_dtype else None + b_internal_type = self.b_smem_alloc_dtype if self.b_smem_alloc_dtype is not self.b_dtype else None def make_tma_maybe_pair( make_atom, @@ -648,17 +566,10 @@ def assign_device_members( ) self.fc1_k_tile_count = ceil_div(bound_hidden_size, self.mma_tile_k) self.fc2_k_tile_count = ceil_div(bound_intermediate_gateup_size // 2, self.mma_tile_k) - self.sfb_vmnk_coord = ( - 0, - cta_coord_in_cluster[0], - cta_coord_in_cluster[1], - cta_coord_in_cluster[2], - ) + self.sfb_vmnk_coord = (0, cta_coord_in_cluster[0], cta_coord_in_cluster[1], cta_coord_in_cluster[2]) @cute.jit - def create_ab_pipeline( - self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer - ) -> pipeline.PipelineAsync: + def create_ab_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> pipeline.PipelineAsync: """Create the shared two-producer TMA-to-UMMA pipeline.""" producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 2) return PipelineTmaUmmaMixedCga.create( @@ -679,9 +590,7 @@ def create_ab_pipeline( def create_acc_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer): """Create the UMMA-to-epilogue accumulator pipeline.""" producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) - consumer_thread_count = ( - self.num_accumulator_consumer_warps_per_cta * 32 * self.mma_cta_count - ) + consumer_thread_count = self.num_accumulator_consumer_warps_per_cta * 32 * self.mma_cta_count consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, consumer_thread_count) return pipeline.PipelineUmmaAsync.create( barrier_storage=smem_workspace.ptr(self.acc_pipeline_mbarriers_region, smem_base), @@ -693,9 +602,7 @@ def create_acc_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Poi ) @cute.jit - def create_tmem_allocator( - self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer, *, allocator_warp_id: int - ): + def create_tmem_allocator(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer, *, allocator_warp_id: int): """Bind the TMEM allocator to finalized SMEM regions.""" allocation_barrier = pipeline.NamedBarrier( barrier_id=self.tmem_allocation_barrier_id, @@ -706,9 +613,7 @@ def create_tmem_allocator( barrier_for_retrieve=allocation_barrier, allocator_warp_id=allocator_warp_id, is_two_cta=self.use_2cta_instrs, - two_cta_tmem_dealloc_mbar_ptr=smem_workspace.ptr( - self.tmem_deallocation_mbarrier_region, smem_base - ), + two_cta_tmem_dealloc_mbar_ptr=smem_workspace.ptr(self.tmem_deallocation_mbarrier_region, smem_base), arch=self.architecture, ) @@ -745,17 +650,10 @@ def run_tma_a( multicast_mask = None if cutlass.const_expr(self.is_a_mcast or self.use_2cta_instrs): multicast_mask = tma_multicast_mask( - self.cluster_shape_mn, - None, - self.cta_coord_in_cluster, - None, - self.use_2cta_instrs, - "a", + self.cluster_shape_mn, None, self.cta_coord_in_cluster, None, self.use_2cta_instrs, "a" ) - a_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_vmnk, (0, 0, None, 0)).shape - ) + a_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_vmnk, (0, 0, None, 0)).shape) sfa_cta_layout = a_cta_layout thread_mma = self.tiled_mma.get_slice(self.mma_cta_index) @@ -765,16 +663,10 @@ def run_tma_a( if is_fc1: iket.range_push("tma_weight_fc1") k_tile_count = self.fc1_k_tile_count - real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor( - "a", fc1_tma_a_tensor, work_tile - ) - real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfa", fc1_tma_sfa_tensor, work_tile - ) + real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor("a", fc1_tma_a_tensor, work_tile) + real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor("sfa", fc1_tma_sfa_tensor, work_tile) - global_a = cute.local_tile( - real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) - ) + global_a = cute.local_tile(real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None)) global_sfa = cute.local_tile( real_sfa, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) ) @@ -829,16 +721,10 @@ def run_tma_a( else: iket.range_push("tma_weight_fc2") k_tile_count = self.fc2_k_tile_count - real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor( - "a", fc2_tma_a_tensor, work_tile - ) - real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfa", fc2_tma_sfa_tensor, work_tile - ) + real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor("a", fc2_tma_a_tensor, work_tile) + real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor("sfa", fc2_tma_sfa_tensor, work_tile) - global_a = cute.local_tile( - real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) - ) + global_a = cute.local_tile(real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None)) global_sfa = cute.local_tile( real_sfa, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) ) @@ -927,15 +813,11 @@ def run_tma_b( fc2_tma_b_partition_atom = fc2_tma_b_atom_or_pair[0] fc2_tma_sfb_partition_atom = fc2_tma_sfb_atom_or_pair[0] - fc1_tma_b_exec_atom = _make_executable_tma_atom_for_cluster( - fc1_tma_b_atom_or_pair, self.is_fallback_cluster - ) + fc1_tma_b_exec_atom = _make_executable_tma_atom_for_cluster(fc1_tma_b_atom_or_pair, self.is_fallback_cluster) fc1_tma_sfb_exec_atom = _make_executable_tma_atom_for_cluster( fc1_tma_sfb_atom_or_pair, self.is_fallback_cluster ) - fc2_tma_b_exec_atom = _make_executable_tma_atom_for_cluster( - fc2_tma_b_atom_or_pair, self.is_fallback_cluster - ) + fc2_tma_b_exec_atom = _make_executable_tma_atom_for_cluster(fc2_tma_b_atom_or_pair, self.is_fallback_cluster) fc2_tma_sfb_exec_atom = _make_executable_tma_atom_for_cluster( fc2_tma_sfb_atom_or_pair, self.is_fallback_cluster ) @@ -960,20 +842,14 @@ def run_tma_b( "sfb", ) - b_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_vmnk, (0, None, 0, 0)).shape - ) - sfb_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape - ) + b_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_vmnk, (0, None, 0, 0)).shape) + sfb_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape) # Rescale fallback CTA coordinates into the preferred tile partition. b_partition_coord = self.main_vmnk_coord[1] sfb_partition_coord = self.sfb_vmnk_coord[1] if cutlass.const_expr(self.is_mixed_cga): if self.is_fallback_cluster: - split_factor = ( - self.cluster_shape_mn[0] // self.resolved_fallback_cluster_shape_mn[0] - ) + split_factor = self.cluster_shape_mn[0] // self.resolved_fallback_cluster_shape_mn[0] b_partition_coord = b_partition_coord * Int32(split_factor) sfb_partition_coord = sfb_partition_coord * Int32(split_factor) thread_mma = self.tiled_mma.get_slice(self.mma_cta_index) @@ -988,24 +864,17 @@ def run_tma_b( kernel_extension.wait_for_input(work_tile) iket.range_pop() k_tile_count = self.fc1_k_tile_count - real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor( - "b", fc1_tma_b_tensor, work_tile - ) - real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfb", fc1_tma_sfb_tensor, work_tile - ) + real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor("b", fc1_tma_b_tensor, work_tile) + real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor("sfb", fc1_tma_sfb_tensor, work_tile) if cutlass.const_expr(self.use_2cta_instrs): if not self.is_leader_cta: load_shift = dynamic_mainloop.compute_non_leader_cta_load_shift( - valid_tokens_in_tile=work_tile.valid_tokens_in_cta_tile, - mma_tiler_n=self.mma_tiler_mnk[1], + valid_tokens_in_tile=work_tile.valid_tokens_in_cta_tile, mma_tiler_n=self.mma_tiler_mnk[1] ) real_b = cute.domain_offset((load_shift, 0, 0), real_b) - global_b = cute.local_tile( - real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None) - ) + global_b = cute.local_tile(real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None)) global_sfb = cute.local_tile( real_sfb, cute.slice_(self.mma_tiler_sfb, (0, None, None)), (None, None, None) ) @@ -1049,11 +918,7 @@ def run_tma_b( tma_desc_ptr=b_descriptor_ptr, mcast_mask=b_multicast_mask, ) - cute.copy( - b_issue_atom, - global_b_slice[(None, k_tile_idx)], - partitioned_smem_b[(None, stage_index)], - ) + cute.copy(b_issue_atom, global_b_slice[(None, k_tile_idx)], partitioned_smem_b[(None, stage_index)]) sfb_issue_atom = bind_executable_tma_load_fields( fc1_tma_sfb_exec_atom, tma_bar_ptr=tma_barrier, @@ -1061,9 +926,7 @@ def run_tma_b( mcast_mask=sfb_multicast_mask, ) cute.copy( - sfb_issue_atom, - global_sfb_slice[(None, k_tile_idx)], - partitioned_smem_sfb[(None, stage_index)], + sfb_issue_atom, global_sfb_slice[(None, k_tile_idx)], partitioned_smem_sfb[(None, stage_index)] ) else: iket.range_push("tma_token_fc2") @@ -1079,24 +942,17 @@ def run_tma_b( iket.range_pop() k_tile_count = self.fc2_k_tile_count - real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor( - "b", fc2_tma_b_tensor, work_tile - ) - real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfb", fc2_tma_sfb_tensor, work_tile - ) + real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor("b", fc2_tma_b_tensor, work_tile) + real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor("sfb", fc2_tma_sfb_tensor, work_tile) if cutlass.const_expr(self.use_2cta_instrs): if not self.is_leader_cta: load_shift = dynamic_mainloop.compute_non_leader_cta_load_shift( - valid_tokens_in_tile=work_tile.valid_tokens_in_cta_tile, - mma_tiler_n=self.mma_tiler_mnk[1], + valid_tokens_in_tile=work_tile.valid_tokens_in_cta_tile, mma_tiler_n=self.mma_tiler_mnk[1] ) real_b = cute.domain_offset((load_shift, 0, 0), real_b) - global_b = cute.local_tile( - real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None) - ) + global_b = cute.local_tile(real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None)) global_sfb = cute.local_tile( real_sfb, cute.slice_(self.mma_tiler_sfb, (0, None, None)), (None, None, None) ) @@ -1140,11 +996,7 @@ def run_tma_b( tma_desc_ptr=b_descriptor_ptr, mcast_mask=b_multicast_mask, ) - cute.copy( - b_issue_atom, - global_b_slice[(None, k_tile_idx)], - partitioned_smem_b[(None, stage_index)], - ) + cute.copy(b_issue_atom, global_b_slice[(None, k_tile_idx)], partitioned_smem_b[(None, stage_index)]) sfb_issue_atom = bind_executable_tma_load_fields( fc2_tma_sfb_exec_atom, tma_bar_ptr=tma_barrier, @@ -1152,9 +1004,7 @@ def run_tma_b( mcast_mask=sfb_multicast_mask, ) cute.copy( - sfb_issue_atom, - global_sfb_slice[(None, k_tile_idx)], - partitioned_smem_sfb[(None, stage_index)], + sfb_issue_atom, global_sfb_slice[(None, k_tile_idx)], partitioned_smem_sfb[(None, stage_index)] ) iket.range_pop() @@ -1178,14 +1028,10 @@ def run_mma( tmem_allocator.wait_for_alloc() accumulator_pointer = tmem_allocator.retrieve_ptr(self.acc_dtype) - accumulator_layout = cute.make_layout( - self.accumulator_shape, stride=self.accumulator_stride - ) + accumulator_layout = cute.make_layout(self.accumulator_shape, stride=self.accumulator_stride) accumulator_base = cute.make_tensor(accumulator_pointer, accumulator_layout) - sfa_pointer = cute.recast_ptr( - accumulator_pointer + self.num_accumulator_tmem_cols, dtype=self.sf_dtype - ) + sfa_pointer = cute.recast_ptr(accumulator_pointer + self.num_accumulator_tmem_cols, dtype=self.sf_dtype) sfa_layout = blockscaled_utils.make_tmem_layout_sfa( self.tiled_mma, self.mma_tiler_mnk, @@ -1195,8 +1041,7 @@ def run_mma( tmem_sfa = cute.make_tensor(sfa_pointer, sfa_layout) sfb_pointer = cute.recast_ptr( - accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, - dtype=self.sf_dtype, + accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, dtype=self.sf_dtype ) sfb_layout = blockscaled_utils.make_tmem_layout_sfb( self.tiled_mma, @@ -1237,10 +1082,7 @@ def run_mma( if cutlass.const_expr(self.cta_tile_n == 64): sfb_shift = (work_tile.tile_n_idx % Int32(2)) * Int32(2) shifted_sfb_pointer = cute.recast_ptr( - accumulator_pointer - + self.num_accumulator_tmem_cols - + self.num_sfa_tmem_cols - + sfb_shift, + accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols + sfb_shift, dtype=self.sf_dtype, ) mma_tmem_sfb = cute.make_tensor(shifted_sfb_pointer, sfb_layout) @@ -1256,12 +1098,8 @@ def run_mma( next_full_status = ab_pipeline.consumer_try_wait(next_ab_pipeline_state) s2t_stage_coord = (None, None, None, None, ab_pipeline_state.index) - cute.copy( - tiled_copy_sfa, partitioned_smem_sfa[s2t_stage_coord], partitioned_tmem_sfa - ) - cute.copy( - tiled_copy_sfb, partitioned_smem_sfb[s2t_stage_coord], partitioned_tmem_sfb - ) + cute.copy(tiled_copy_sfa, partitioned_smem_sfa[s2t_stage_coord], partitioned_tmem_sfa) + cute.copy(tiled_copy_sfb, partitioned_smem_sfb[s2t_stage_coord], partitioned_tmem_sfb) stage_coord = (None, None, None, ab_pipeline_state.index) dynamic_mainloop.issue_dynamic_block_scaled_mma_tile( quant_kind=self.quant_kind, diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py index e5482f1c4491..7dad5c86e02b 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py @@ -1,11 +1,11 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Fused pull-dispatch MegaMoE kernel composition for Blackwell.""" +import hashlib + from typing import ClassVar, Optional, Tuple import cuda.bindings.driver as cuda + import cutlass import cutlass.cute as cute import cutlass.pipeline as pipeline @@ -14,23 +14,20 @@ from cutlass.cute.nvgpu import OperandMajorMode from cutlass.cutlass_dsl import Boolean, Int32 -from .....api import ( - ImplDesc, - KernelClass, - OptionalRequirement, - ProblemDesc, - StaticOrRuntimeIntegerType, -) +from .....api import ImplDesc, KernelClass, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType from .....communication.nvlink_domain.token_comm import TokenCommArgs, TokenCommNonDeterministic -from .....helpers.cute_py_helpers import ( - Tcgen05MmaInstruction, - make_tcgen05_tmem_plan, - tcgen05_block_scaled_acc_dtype, -) +from .....helpers.cute_py_helpers import Tcgen05MmaInstruction, make_tcgen05_tmem_plan, tcgen05_block_scaled_acc_dtype from .....helpers.device_workspace import DeviceWorkspace +from .....helpers.dlb_expert_count import resolve_dlb_expert_counts from .....helpers.iket_compat import iket +from .....helpers.megamoe_aot import make_megamoe_aot_callable from .....helpers.smem_workspace import SmemWorkspace +from .....helpers.topk_reduce import TopkReduce from .....helpers.utils import ceil_div, round_up +from .....helpers.helper_weight_ready_gate import ( + derive_helper_weight_ready_topology, + make_helper_weight_ready_generation_gate, +) from .....quant_def import CombineFormat, QuantKind from ....schedulers.base import WorkIdAcquisitionMode from ....schedulers.fc12_scheduler import ( @@ -39,15 +36,11 @@ minimum_phase_interleave_hint, ) from ....schedulers.non_clc_mixed_cga import NonClcMixedCgaConfig -from ...custom_mix_cga_helpers import ( - TmaAtomOrPair, - pipeline_init_arrive_mixed_cga, - pipeline_init_wait_mixed_cga, -) +from ...custom_mix_cga_helpers import TmaAtomOrPair, pipeline_init_arrive_mixed_cga, pipeline_init_wait_mixed_cga from .block_scaled_swap_ab_fc12_epilogue import GatedActEpilogueArgs, SwapABGatedActEpilogue from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension from .block_scaled_swap_ab_fc12_mainloop import BlockScaledSwapAbFc12Mainloop -from .topk_reduce import TopkReduce + _aot_symbol_prefix = "blackwell_mega_moe_aot" @@ -110,10 +103,18 @@ def impl_desc_require(cls) -> dict[str, type]: "token_in_flag_batch": int, "token_back_mode": str, "reduce_topk_in_kernel": bool, + "helper_expert_count": OptionalRequirement(int), + "total_helper_slots": OptionalRequirement(int), } def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self._validate_desc_inputs(problem_desc, impl_desc) + problem_desc, impl_desc = resolve_dlb_expert_counts( + problem_desc, impl_desc, world_size=problem_desc["world_size"] + ) + # The Blackwell epilogue publishes expert counters, not token-tile counters. + if impl_desc.get("token_back_ready_granularity", "expert") != "expert": + raise ValueError("Blackwell requires token_back_ready_granularity='expert'.") self.expert_count = problem_desc["expert_count"] self.intermediate_gateup_size = problem_desc["intermediate_gateup_size"] @@ -154,10 +155,23 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.token_in_flag_batch = impl_desc["token_in_flag_batch"] self.token_back_mode = impl_desc["token_back_mode"] self.reduce_topk_in_kernel = impl_desc["reduce_topk_in_kernel"] - + self.helper_expert_count = impl_desc.get("helper_expert_count", 0) + if type(self.helper_expert_count) is not int or self.helper_expert_count < 0: + raise ValueError("helper_expert_count must be a non-negative exact int") self.occupancy = 1 self.architecture = "sm_100" self.local_expert_count = self.expert_count // self.world_size + if self.helper_expert_count > 0: + helper_topology = derive_helper_weight_ready_topology( + self.local_expert_count, + self.helper_expert_count, + ) + self.home_expert_count = helper_topology.home_expert_count + else: + # Preserve the baseline's runtime-expert-count specialization. A + # static M is required only when an S>0 helper boundary must be + # compiled into the scheduler gate. + self.home_expert_count = self.local_expert_count self.threads_per_cta = 16 * 32 if self.token_back_mode == "standalone_warps" else 12 * 32 # nosec B105 self.other_warp_register_count = 64 if self.token_back_mode == "standalone_warps" else 72 # nosec B105 @@ -172,9 +186,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.preferred_cluster_count = self.mixed_cga_config.preferred_cluster_count self.fallback_cluster_count = self.mixed_cga_config.fallback_cluster_count self.resolved_fallback_cluster_shape_mn = ( - self.cluster_shape_mn - if self.fallback_cluster_shape_mn is None - else self.fallback_cluster_shape_mn + self.cluster_shape_mn if self.fallback_cluster_shape_mn is None else self.fallback_cluster_shape_mn ) if self.launch_cluster_count != self.mixed_cga_config.launch_cluster_cnt_merge_as_preferred: raise ValueError( @@ -189,11 +201,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: a_type=self.a_dtype, b_type=self.b_dtype, acc_type=self.acc_dtype, - instruction_mnk=( - self.mma_tiler_mnk[0], - self.mma_tiler_mnk[1], - self.quant_kind.instruction_k("1x"), - ), + instruction_mnk=(self.mma_tiler_mnk[0], self.mma_tiler_mnk[1], self.quant_kind.instruction_k("1x")), participates=mma_cta_count, sfa_type=self.sf_dtype, sfb_type=self.sf_dtype, @@ -201,9 +209,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: ) tmem_plan = make_tcgen05_tmem_plan(mma_instruction, self.architecture, self.mma_tiler_mnk) tokens_per_fc1_ready_slot = self.mma_tiler_mnk[1] * self.cluster_shape_mn[1] - hidden_per_fc2_cluster_tile = ( - self.mma_tiler_mnk[0] // mma_cta_count * self.cluster_shape_mn[0] - ) + hidden_per_fc2_cluster_tile = self.mma_tiler_mnk[0] // mma_cta_count * self.cluster_shape_mn[0] fc2_done_signals_per_token_tile = ( ceil_div(self.hidden_size, hidden_per_fc2_cluster_tile) * self.cluster_shape_mn[0] @@ -218,9 +224,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: "tokens_per_fc1_ready_slot": tokens_per_fc1_ready_slot, "fc2_done_signals_per_token_tile": fc2_done_signals_per_token_tile, "promised_launchable_sm_count": self.mixed_cga_config.total_cta_cnt, - "token_back_schedule_mode": ( - "atomic_counter" if self.work_id_mode == "atomic_counter" else "static" - ), + "token_back_schedule_mode": ("atomic_counter" if self.work_id_mode == "atomic_counter" else "static"), } ) self.token_comm = TokenCommNonDeterministic(problem_desc, token_comm_impl_desc) @@ -246,9 +250,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: } ) scheduler_type = ( - BlackwellFusedFc12Scheduler - if self.schedule_mode == "grouped" - else PhaseInterleavedFc12Scheduler + BlackwellFusedFc12Scheduler if self.schedule_mode == "grouped" else PhaseInterleavedFc12Scheduler ) self.scheduler = scheduler_type(fc12_problem_desc, resolved_impl_desc) self.epilogue = SwapABGatedActEpilogue(fc12_problem_desc, resolved_impl_desc) @@ -259,13 +261,9 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: if self.epilogue.sf_vec_size != self.sf_vec_size: raise ValueError("Epilogue and Mainloop scale vector sizes must match.") self._device_workspace = self._build_device_workspace() - self._mainloop, self._smem_workspace = self._build_mainloop_and_smem( - fc12_problem_desc, resolved_impl_desc - ) + self._mainloop, self._smem_workspace = self._build_mainloop_and_smem(fc12_problem_desc, resolved_impl_desc) self._topk_reduce = ( - None - if self.reduce_topk_in_kernel - else TopkReduce(self.hidden_size, self.topk, self.combine_format) + None if self.reduce_topk_in_kernel else TopkReduce(self.hidden_size, self.topk, self.combine_format) ) def _validate_geometry(self) -> None: @@ -277,9 +275,7 @@ def _validate_geometry(self) -> None: if any(static_expert_dimensions) and not all(static_expert_dimensions): raise ValueError("MegaMoE expert dimensions must be either all static or all runtime.") if not all(static_expert_dimensions): - raise NotImplementedError( - "The MegaMoE Kernel currently requires static expert dimensions." - ) + raise NotImplementedError("The MegaMoE Kernel currently requires static expert dimensions.") if self.expert_count <= 0 or self.intermediate_gateup_size <= 0 or self.hidden_size <= 0: raise ValueError("MegaMoE expert dimensions must be positive.") if self.world_size <= 0 or self.expert_count % self.world_size != 0: @@ -287,9 +283,7 @@ def _validate_geometry(self) -> None: if self.topk <= 0 or self.topk > self.expert_count: raise ValueError("topk must be positive and no greater than expert_count.") if self.topk_index_dtype not in (cutlass.Int32, cutlass.Int64): - raise ValueError( - f"topk_index_dtype must be Int32 or Int64, got {self.topk_index_dtype}." - ) + raise ValueError(f"topk_index_dtype must be Int32 or Int64, got {self.topk_index_dtype}.") if self.reduce_topk_in_kernel and not self.apply_topk_at_fc1: raise ValueError("In-kernel top-k reduction requires apply_topk_at_fc1=True.") if self.intermediate_gateup_size % 2 != 0: @@ -306,12 +300,8 @@ def _validate_geometry(self) -> None: raise ValueError( f"schedule_policy mode must be 'grouped' or 'phase_interleave', got {self.schedule_mode!r}." ) - if self.hint is not None and ( - isinstance(self.hint, bool) or not isinstance(self.hint, int) or self.hint <= 0 - ): - raise ValueError( - f"schedule_policy hint must be a positive Python int, got {self.hint!r}." - ) + if self.hint is not None and (isinstance(self.hint, bool) or not isinstance(self.hint, int) or self.hint <= 0): + raise ValueError(f"schedule_policy hint must be a positive Python int, got {self.hint!r}.") if self.schedule_mode == "phase_interleave" and self.work_id_mode != "atomic_counter": raise ValueError("phase_interleave currently requires work_id_mode='atomic_counter'.") if len(self.epi_flag_batches) != 2: @@ -345,9 +335,7 @@ def _validate_geometry(self) -> None: if mma_k % (self.sf_vec_size * 4) != 0: raise ValueError("mma_tiler K must be divisible by four scale-factor vectors.") if self.intermediate_gateup_size % (self.sf_vec_size * 4) != 0: - raise ValueError( - "The intermediate dimension must be divisible by four scale-factor vectors." - ) + raise ValueError("The intermediate dimension must be divisible by four scale-factor vectors.") if cluster_n != 1: raise ValueError(f"The swap-AB FC12 path requires cluster N=1, got {cluster_n}.") if cluster_m <= 0 or cluster_m > 16 or cluster_m & (cluster_m - 1): @@ -357,14 +345,8 @@ def _validate_geometry(self) -> None: if self.fallback_cluster_shape_mn is not None: fallback_cluster_m, fallback_cluster_n = self.fallback_cluster_shape_mn if fallback_cluster_n != 1: - raise ValueError( - f"The swap-AB FC12 fallback path requires cluster N=1, got {fallback_cluster_n}." - ) - if ( - fallback_cluster_m <= 0 - or fallback_cluster_m > 16 - or fallback_cluster_m & (fallback_cluster_m - 1) - ): + raise ValueError(f"The swap-AB FC12 fallback path requires cluster N=1, got {fallback_cluster_n}.") + if fallback_cluster_m <= 0 or fallback_cluster_m > 16 or fallback_cluster_m & (fallback_cluster_m - 1): raise ValueError("fallback cluster M must be a power of two no greater than 16.") if self.use_2cta_instrs and fallback_cluster_m % 2 != 0: raise ValueError("Two-CTA MMA requires an even fallback cluster M.") @@ -462,6 +444,7 @@ def dtype_name(dtype: type) -> str: f"{dtype_name(self.a_dtype)}_{dtype_name(self.b_dtype)}_{dtype_name(self.acc_dtype)}_" f"sfvec{self.sf_vec_size}_a{self.a_major_mode.name.lower()}_b{self.b_major_mode.name.lower()}_" f"e{self.expert_count}_ep{self.world_size}_topk{self.topk}_" + f"helpers{self.helper_expert_count}_" f"topkidx{dtype_name(self.topk_index_dtype)}_" f"h{self.hidden_size}_i{self.intermediate_gateup_size}_maxtoken{self.max_tokens_per_rank}_" f"tile{tile}_cluster{cluster}_{'2cta' if self.use_2cta_instrs else '1cta'}_" @@ -479,6 +462,7 @@ def dtype_name(dtype: type) -> str: f"situ{self.situ_beta}x{self.situ_linear_beta}_" f"{'apply_topk_fc1' if self.apply_topk_at_fc1 else 'apply_topk_fc2'}_" f"{'inkernel_redg' if self.reduce_topk_in_kernel else 'separate_reduce'}" + "_dlb_v1" ) def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): @@ -500,14 +484,10 @@ def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): extents = tuple( - sym_int64(divisibility=math.gcd(int(extent), 128)) - if axis in dynamic_axes - else int(extent) + sym_int64(divisibility=math.gcd(int(extent), 128)) if axis in dynamic_axes else int(extent) for axis, extent in enumerate(shape) ) - return make_fake_compact_tensor( - dtype, extents, stride_order=stride_order, assumed_align=alignment - ) + return make_fake_compact_tensor(dtype, extents, stride_order=stride_order, assumed_align=alignment) tokens = self.max_tokens_per_rank hidden = self.hidden_size @@ -516,20 +496,14 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): experts = self.local_expert_count sf_vec_size = self.sf_vec_size # Weight SF is atom-swizzled and opaque, so its extent is whatever the caller staged, not a kernel choice. - fc1_weight_sf_columns = round_up(intermediate_gateup, 128) * round_up( - hidden // sf_vec_size, 4 - ) - fc2_weight_sf_columns = round_up(hidden, 128) * round_up( - intermediate_downproj // sf_vec_size, 4 - ) + fc1_weight_sf_columns = round_up(intermediate_gateup, 128) * round_up(hidden // sf_vec_size, 4) + fc2_weight_sf_columns = round_up(hidden, 128) * round_up(intermediate_downproj // sf_vec_size, 4) # TopkReduce and the in-kernel REDG both take their output dtype from this tensor. output_dtype = cutlass.BFloat16 fake_arguments = dict( # Swap-AB puts the weights on the MMA A operand, so the token wire dtype is TokenComm's, not a_dtype. - activation=fake_tensor( - self.token_comm.activation_dtype, (tokens, hidden), (1, 0), {0}, 16 - ), + activation=fake_tensor(self.token_comm.activation_dtype, (tokens, hidden), (1, 0), {0}, 16), activation_sf=fake_tensor( self.token_comm.activation_sf_dtype, (tokens, self.token_comm.activation_sf_hidden_padded), @@ -539,18 +513,10 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): ), topk_indices=fake_tensor(self.topk_index_dtype, (tokens, self.topk), (1, 0), {0}, 16), topk_scores=fake_tensor(cutlass.Float32, (tokens, self.topk), (1, 0), {0}, 4), - fc1_weight=fake_tensor( - self.a_dtype, (experts, hidden, intermediate_gateup), (2, 0, 1), {0, 2}, 16 - ), - fc1_weight_sf=fake_tensor( - self.sf_dtype, (experts, fc1_weight_sf_columns), (1, 0), {0}, 16 - ), - fc2_weight=fake_tensor( - self.a_dtype, (experts, intermediate_downproj, hidden), (2, 0, 1), {0, 2}, 16 - ), - fc2_weight_sf=fake_tensor( - self.sf_dtype, (experts, fc2_weight_sf_columns), (1, 0), {0}, 16 - ), + fc1_weight=fake_tensor(self.a_dtype, (experts, hidden, intermediate_gateup), (2, 0, 1), {0, 2}, 16), + fc1_weight_sf=fake_tensor(self.sf_dtype, (experts, fc1_weight_sf_columns), (1, 0), {0}, 16), + fc2_weight=fake_tensor(self.a_dtype, (experts, intermediate_downproj, hidden), (2, 0, 1), {0, 2}, 16), + fc2_weight_sf=fake_tensor(self.sf_dtype, (experts, fc2_weight_sf_columns), (1, 0), {0}, 16), output_activation=fake_tensor(output_dtype, (tokens, hidden), (1, 0), {0}, 16), # Opaque byte workspaces: a bare base pointer, like to_cute_ptr. The 128-byte promise matches the # alignment the registered regions assume, so the caller must hand over 128-byte-aligned bases. @@ -565,6 +531,16 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): ), stream=make_fake_stream(), ) + if self.helper_expert_count > 0: + fake_arguments.update( + hot_expert_weight_ready_flags=make_ptr( + cutlass.Uint64, + 0, + AddressSpace.gmem, + assumed_align=8, + ), + hot_expert_weight_ready_generation=cutlass.Uint64(1), + ) if self.quant_kind.uses_global_scale: # nvfp4 carries per-expert dequant scalars. Under an e8m0 scale they do not exist at # all, and declaring them here would put three tensors in the ABI that the caller has @@ -579,15 +555,21 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): if out_path is None: return compiled compiled.export_to_c( - out_path, function_name=_aot_symbol_prefix, export_only_tvm_ffi_symbols=True + out_path, + function_name=self._aot_symbol_name(), + export_only_tvm_ffi_symbols=True, ) return out_path - @staticmethod - def load_compiled(path: str): + def _aot_symbol_name(self) -> str: + identity = hashlib.sha256(self.name().encode()).hexdigest()[:16] + return f"{_aot_symbol_prefix}_{identity}" + + def load_compiled(self, path: str): from cutlass.cute.runtime import load_module - return load_module(path, enable_tvm_ffi=True)[_aot_symbol_prefix] + function = getattr(load_module(path, enable_tvm_ffi=True), self._aot_symbol_name()) + return make_megamoe_aot_callable(function, self) def _build_mainloop_and_smem( self, problem_desc: ProblemDesc, resolved_impl_desc: ImplDesc @@ -603,9 +585,7 @@ def _build_mainloop_and_smem( if mainloop_smem_budget_bytes <= 0: raise ValueError("Fixed kernel components leave no SMEM budget for the Mainloop.") - mainloop_impl_desc = ImplDesc( - {**resolved_impl_desc, "mainloop_smem_budget_bytes": mainloop_smem_budget_bytes} - ) + mainloop_impl_desc = ImplDesc({**resolved_impl_desc, "mainloop_smem_budget_bytes": mainloop_smem_budget_bytes}) mainloop = BlockScaledSwapAbFc12Mainloop(problem_desc, mainloop_impl_desc) mainloop.register_smem_regions(smem_workspace) @@ -631,16 +611,34 @@ def __call__( fc1_alpha: Optional[cute.Tensor] = None, # (local_experts,) fc2_alpha: Optional[cute.Tensor] = None, # (local_experts,) fc1_norm_const: Optional[cute.Tensor] = None, # (local_experts,) + hot_expert_weight_ready_flags: Optional[cute.Pointer] = None, # uint64[EP] + hot_expert_weight_ready_generation: Optional[cutlass.Uint64] = None, ) -> None: """Launch router, fused MegaMoE compute, and optional TopK reduce.""" def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: return cute.make_tensor(tensor.iterator, cute.make_layout(shape, stride=tensor.stride)) - if cutlass.const_expr(topk_indices.element_type is not self.topk_index_dtype): - raise TypeError( - f"topk_indices must be the declared {self.topk_index_dtype}, got {topk_indices.dtype}." + helper_ready_bundle = ( + hot_expert_weight_ready_flags, + hot_expert_weight_ready_generation, + ) + if cutlass.const_expr( + ( + self.helper_expert_count > 0 + and not all(value is not None for value in helper_ready_bundle) + ) + or ( + self.helper_expert_count == 0 + and any(value is not None for value in helper_ready_bundle) + ) + ): + raise ValueError( + "helper READY flags/generation must both be present iff helper_expert_count > 0" ) + + if cutlass.const_expr(topk_indices.element_type is not self.topk_index_dtype): + raise TypeError(f"topk_indices must be the declared {self.topk_index_dtype}, got {topk_indices.dtype}.") for operand_name, operand, expected_dtype in ( ("activation", activation, self.b_dtype), ("activation_sf", activation_sf, self.sf_dtype), @@ -729,19 +727,14 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: padded_intermediate_gateup = cute.round_up(intermediate_gateup, sf_vec_size * 4) expected_fc1_weight_sf_columns = padded_intermediate_gateup * padded_hidden // sf_vec_size if cutlass.const_expr( - isinstance(fc1_weight_sf.shape[1], int) - and isinstance(expected_fc1_weight_sf_columns, int) + isinstance(fc1_weight_sf.shape[1], int) and isinstance(expected_fc1_weight_sf_columns, int) ): if cutlass.const_expr(fc1_weight_sf.shape[1] != expected_fc1_weight_sf_columns): raise ValueError("fc1_weight_sf has an incompatible column count.") fc1_sfa = cute.make_tensor( fc1_weight_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - cutlass.Int32(padded_intermediate_gateup), - cutlass.Int32(padded_hidden), - cutlass.Int32(experts), - ), + (cutlass.Int32(padded_intermediate_gateup), cutlass.Int32(padded_hidden), cutlass.Int32(experts)), sf_vec_size, ), ) @@ -759,31 +752,21 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: padded_fc2_intermediate = cute.round_up(intermediate_fc2, sf_vec_size * 4) expected_fc2_weight_sf_columns = padded_fc2_hidden * padded_fc2_intermediate // sf_vec_size if cutlass.const_expr( - isinstance(fc2_weight_sf.shape[1], int) - and isinstance(expected_fc2_weight_sf_columns, int) + isinstance(fc2_weight_sf.shape[1], int) and isinstance(expected_fc2_weight_sf_columns, int) ): if cutlass.const_expr(fc2_weight_sf.shape[1] != expected_fc2_weight_sf_columns): raise ValueError("fc2_weight_sf has an incompatible column count.") fc2_sfa = cute.make_tensor( fc2_weight_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - cutlass.Int32(padded_fc2_hidden), - cutlass.Int32(padded_fc2_intermediate), - cutlass.Int32(experts_fc2), - ), + (cutlass.Int32(padded_fc2_hidden), cutlass.Int32(padded_fc2_intermediate), cutlass.Int32(experts_fc2)), sf_vec_size, ), ) fc2_sfb = cute.make_tensor( fc1_output_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - fc1_output_sf.shape[0], - cutlass.Int32(fc1_output_sf.shape[1] * sf_vec_size), - singleton, - ), - sf_vec_size, + (fc1_output_sf.shape[0], cutlass.Int32(fc1_output_sf.shape[1] * sf_vec_size), singleton), sf_vec_size ), ) if cutlass.const_expr(self.reduce_topk_in_kernel): @@ -791,28 +774,19 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: output_activation.iterator, cute.make_layout( (self.max_tokens_per_rank, 1, hidden), - stride=( - output_activation.stride[0], - output_activation.stride[0], - output_activation.stride[1], - ), + stride=(output_activation.stride[0], output_activation.stride[0], output_activation.stride[1]), ), ) pre_reduced_activation_sf = None else: - pre_reduced_activation = self.token_comm.pre_reduced_activation_tensor( - self._device_workspace - ) - pre_reduced_activation_sf = self.token_comm.pre_reduced_activation_sf_tensor( - self._device_workspace - ) + pre_reduced_activation = self.token_comm.pre_reduced_activation_tensor(self._device_workspace) + pre_reduced_activation_sf = self.token_comm.pre_reduced_activation_sf_tensor(self._device_workspace) if cutlass.const_expr(self.token_comm.token_back_push_data): fc2_output = self.token_comm.fc2_activation_tensor(self._device_workspace) else: fc2_output = rewrite_tensor_shape( - pre_reduced_activation, - (pre_reduced_activation.shape[0], pre_reduced_activation.shape[1], hidden), + pre_reduced_activation, (pre_reduced_activation.shape[0], pre_reduced_activation.shape[1], hidden) ) fc2_output_sf = self.token_comm.fc2_activation_sf_tensor(self._device_workspace) @@ -886,6 +860,8 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: fc1_alpha, fc2_alpha, fc1_norm_const, + hot_expert_weight_ready_flags, + hot_expert_weight_ready_generation, ) if cutlass.const_expr(self.mixed_cga_config.is_mixed): kernel.launch( @@ -909,11 +885,7 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: if cutlass.const_expr(not self.reduce_topk_in_kernel): reduce_scores = None if cutlass.const_expr(self.apply_topk_at_fc1) else topk_scores self._topk_reduce( - pre_reduced_activation, - pre_reduced_activation_sf, - output_activation, - reduce_scores, - stream, + pre_reduced_activation, pre_reduced_activation_sf, output_activation, reduce_scores, stream ) @cute.kernel @@ -953,6 +925,8 @@ def _kernel( fc1_alpha: Optional[cute.Tensor], fc2_alpha: Optional[cute.Tensor], fc1_norm_const: Optional[cute.Tensor], + hot_expert_weight_ready_flags: Optional[cute.Pointer], + hot_expert_weight_ready_generation: Optional[cutlass.Uint64], ): """Compose TokenComm, Scheduler, Mainloop, and Epilogue.""" self._mainloop.materialize_codegen_members() @@ -974,8 +948,7 @@ def _kernel( fc1_output_sf = cute.make_tensor( fc1_output_sf_storage.iterator, blockscaled_utils.tile_atom_to_shape_SF( - (fc1_output_sf_storage.shape[0], intermediate_gateup // 2, cutlass.Int32(1)), - self.epilogue.sf_vec_size, + (fc1_output_sf_storage.shape[0], intermediate_gateup // 2, cutlass.Int32(1)), self.epilogue.sf_vec_size ), ) @@ -986,9 +959,7 @@ def _kernel( is_fallback_cluster = Boolean(False) if cutlass.const_expr(self.mixed_cga_config.is_mixed): active_cluster_m, _, _ = cute.arch.block_in_cluster_dim() - is_fallback_cluster = active_cluster_m == Int32( - self.resolved_fallback_cluster_shape_mn[0] - ) + is_fallback_cluster = active_cluster_m == Int32(self.resolved_fallback_cluster_shape_mn[0]) cta_rank_in_cluster = cute.arch.block_idx_in_cluster() cta_coord_in_cluster = (cta_rank_in_cluster, Int32(0), Int32(0)) linear_cta_idx = block_idx[0] + block_idx[2] * Int32(self.cluster_shape_mn[0]) @@ -1012,12 +983,7 @@ def _kernel( # Bind the active cluster before deriving the AB pipeline arrival count. self._mainloop.assign_device_members( - self._smem_workspace, - smem_base, - cta_coord_in_cluster, - is_fallback_cluster, - hidden, - intermediate_gateup, + self._smem_workspace, smem_base, cta_coord_in_cluster, is_fallback_cluster, hidden, intermediate_gateup ) ab_pipeline = self._mainloop.create_ab_pipeline(self._smem_workspace, smem_base) tma_a_pipeline_state = pipeline.make_pipeline_state( @@ -1045,39 +1011,34 @@ def _kernel( ) scheduler = self.scheduler - fc2_spin_threshold = ( - intermediate_gateup + self._mainloop.cta_tile_m - 1 - ) // self._mainloop.cta_tile_m + fc2_spin_threshold = (intermediate_gateup + self._mainloop.cta_tile_m - 1) // self._mainloop.cta_tile_m + helper_weight_ready_gate = None + if cutlass.const_expr(self.helper_expert_count > 0): + helper_weight_ready_gate = make_helper_weight_ready_generation_gate( + memory_slot_count=self.local_expert_count, + helper_count=self.helper_expert_count, + source_count=self.world_size, + terminal_flags=hot_expert_weight_ready_flags, + expected_generation=hot_expert_weight_ready_generation, + ) kernel_extension = BlockScaledSwapAbFc12Extension( sf_vec_size=self.sf_vec_size, fc1_done_counter_pointer=fc1_done_counter.iterator, fc2_spin_threshold=fc2_spin_threshold, - fc1_ready_counter_pointer=self.token_comm.fc1_ready_counter_pointer( - self._device_workspace - ), + fc1_ready_counter_pointer=self.token_comm.fc1_ready_counter_pointer(self._device_workspace), ) optional_epilogue_args = GatedActEpilogueArgs( - fc1_alpha=fc1_alpha, - fc2_alpha=fc2_alpha, - fc1_norm_const=fc1_norm_const, - topk_scores=pool_topk_scores, + fc1_alpha=fc1_alpha, fc2_alpha=fc2_alpha, fc1_norm_const=fc1_norm_const, topk_scores=pool_topk_scores ) pipeline_init_arrive_mixed_cga( - self.cluster_shape_mn, - self.resolved_fallback_cluster_shape_mn, - is_fallback_cluster, - is_relaxed=True, + self.cluster_shape_mn, self.resolved_fallback_cluster_shape_mn, is_fallback_cluster, is_relaxed=True ) pipeline_init_wait_mixed_cga( self.cluster_shape_mn, self.resolved_fallback_cluster_shape_mn, is_fallback_cluster ) - finalize_barrier_thread_count = ( - 12 * 32 if self.token_back_mode == "standalone_warps" else 8 * 32 # nosec B105 - ) - finalize_barrier = pipeline.NamedBarrier( - barrier_id=13, num_threads=finalize_barrier_thread_count - ) + finalize_barrier_thread_count = 12 * 32 if self.token_back_mode == "standalone_warps" else 8 * 32 # nosec B105 + finalize_barrier = pipeline.NamedBarrier(barrier_id=13, num_threads=finalize_barrier_thread_count) if warp_idx == self.scheduler_warp_id: cute.arch.warpgroup_reg_dealloc(self.scheduler_register_count) @@ -1091,9 +1052,15 @@ def _kernel( iket.range_push("scheduler.gen_work") work_tile = scheduler.gen_next_work() iket.range_pop() + helper_ready_latched = Int32(0) while work_tile.is_valid_tile: iket.range_push("scheduler.publish_work") - scheduler.publish_work(kernel_extension.prepare_work_tile(work_tile)) + prepared_tile = kernel_extension.prepare_work_tile(work_tile) + if cutlass.const_expr(helper_weight_ready_gate is not None): + helper_ready_latched = helper_weight_ready_gate.wait_before_first_helper_publish( + prepared_tile.expert_idx, helper_ready_latched + ) + scheduler.publish_work(prepared_tile) iket.range_pop() iket.range_push("scheduler.gen_work") work_tile = scheduler.gen_next_work() @@ -1188,9 +1155,7 @@ def _kernel( optional_epilogue_args, ) tmem_allocator.relinquish_alloc_permit() - tmem_allocator.free( - tmem_allocator.retrieve_ptr(self.acc_dtype), self._mainloop.num_tmem_alloc_cols - ) + tmem_allocator.free(tmem_allocator.retrieve_ptr(self.acc_dtype), self._mainloop.num_tmem_alloc_cols) iket.range_pop() finalize_barrier.arrive_unaligned() @@ -1201,8 +1166,7 @@ def _kernel( self.token_comm.token_in(self._smem_workspace, smem_base) iket.range_pop() if cutlass.const_expr( - self.token_comm.token_back_enabled - and self.token_back_mode != "standalone_warps" # nosec B105 + self.token_comm.token_back_enabled and self.token_back_mode != "standalone_warps" # nosec B105 ): iket.range_push("mega.token_back") self.token_comm.token_back(self._smem_workspace, smem_base) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/dynamic_mainloop.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/dynamic_mainloop.py index b7fedaa598a1..6216b50fb00d 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/dynamic_mainloop.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/dynamic_mainloop.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - # Dynamic UMMA N for fused fc1+fc2 swap-AB MegaMoE kernel. # # Emits the block-scaled UMMA instruction as literal PTX text (via @@ -34,13 +31,16 @@ from typing import Optional +import cutlass import cutlass.cute as cute +from cutlass.cutlass_dsl import dsl_user_op, Int32, Boolean from cutlass._mlir import ir -from cutlass._mlir.dialects import builtin, llvm -from cutlass.cutlass_dsl import Boolean, Int32, dsl_user_op +from cutlass._mlir.dialects import builtin +from cutlass._mlir.dialects import llvm from .....quant_def import QuantKind + # ============================================================================= # Alignment policy (single source of truth) # ============================================================================= @@ -223,20 +223,10 @@ def _tcgen05_mma_block_scaled( """ assert cta_group in (1, 2), f"cta_group must be 1 or 2, got {cta_group}" assert umma_kind in ("mxf4nvf4", "mxf4", "mxf8f6f4"), f"Unsupported UMMA kind {umma_kind!r}" - assert scale_vec_suffix in ("", ".block16", ".block32"), ( - f"Unsupported suffix {scale_vec_suffix!r}" - ) + assert scale_vec_suffix in ("", ".block16", ".block32"), f"Unsupported suffix {scale_vec_suffix!r}" llvm.inline_asm( None, - [ - d_tmem_i32, - a_desc_i64, - b_desc_i64, - idesc_i32, - enable_input_d_i32, - sfa_tmem_i32, - sfb_tmem_i32, - ], + [d_tmem_i32, a_desc_i64, b_desc_i64, idesc_i32, enable_input_d_i32, sfa_tmem_i32, sfb_tmem_i32], "{\n\t" ".reg .pred p;\n\t" "setp.ne.b32 p, $4, 0;\n\t" diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/function_mapping.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/function_mapping.py index c222690227be..83cacc303661 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/function_mapping.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/function_mapping.py @@ -23,13 +23,17 @@ def __post_init__(self) -> None: names = tuple(self.names) sizes = tuple(self.sizes) if not names or len(names) != len(sizes): - raise FunctionMappingError(f"Coordinate-space rank mismatch: {names!r}, {sizes!r}.") + raise FunctionMappingError( + f"Coordinate-space rank mismatch: {names!r}, {sizes!r}." + ) if any(not isinstance(name, str) or not name for name in names): raise FunctionMappingError("Axis names must be non-empty strings.") if len(set(names)) != len(names): raise FunctionMappingError(f"Axis names must be unique: {names!r}.") if any(not isinstance(size, int) or size <= 0 for size in sizes): - raise FunctionMappingError(f"Axis sizes must be positive Python ints: {sizes!r}.") + raise FunctionMappingError( + f"Axis sizes must be positive Python ints: {sizes!r}." + ) object.__setattr__(self, "names", names) object.__setattr__(self, "sizes", sizes) @@ -45,7 +49,9 @@ def axis_size(self, name: str) -> int: def delinearize(self, linear_index: int) -> tuple[int, ...]: if linear_index < 0 or linear_index >= self.size: - raise FunctionMappingError(f"Linear index {linear_index} is outside [0, {self.size}).") + raise FunctionMappingError( + f"Linear index {linear_index} is outside [0, {self.size})." + ) remaining = linear_index coordinate = [] for size in self.sizes: @@ -92,7 +98,9 @@ def _validate_signature(self) -> None: or parameter.default is not inspect.Parameter.empty ] if unsupported: - raise FunctionMappingError(f"Unsupported mapping parameters {unsupported!r}.") + raise FunctionMappingError( + f"Unsupported mapping parameters {unsupported!r}." + ) if set(parameters) != set(self.domain.names): raise FunctionMappingError( f"Mapping parameters {tuple(parameters)!r} must match " @@ -111,14 +119,17 @@ def _normalize_result( f"Mapping result keys {tuple(result)!r} must match " f"codomain axes {self.codomain.names!r}." ) - coordinate = {name: result[name] for name in self.codomain.names} + coordinate = { + name: result[name] for name in self.codomain.names + } elif isinstance(result, Sequence) and not isinstance( result, (str, bytes), ): if len(result) != len(self.codomain.names): raise FunctionMappingError( - f"Mapping result rank {len(result)} must equal {len(self.codomain.names)}." + f"Mapping result rank {len(result)} must equal " + f"{len(self.codomain.names)}." ) coordinate = dict(zip(self.codomain.names, result)) elif len(self.codomain.names) == 1: diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/__init__.py index 52a7a9daf028..97f3397bcf5a 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/__init__.py @@ -1,2 +1,6 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Rubin Local MegaMoE inference kernel.""" + +from .block_scaled_swap_ab_local_mega_moe_kernel import BlockScaledSwapAbLocalMegaMoeKernel + + +__all__ = ["BlockScaledSwapAbLocalMegaMoeKernel"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/custom_mix_cga_helpers.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/custom_mix_cga_helpers.py new file mode 100644 index 000000000000..dececcded527 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/custom_mix_cga_helpers.py @@ -0,0 +1,25 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Rubin source-copy shim for shared mixed-CGA helpers.""" + +# <<>> +from ..blackwell.custom_mix_cga_helpers import ( + PipelineTmaUmmaMixedCga, + TmaAtomOrPair, + bind_executable_tma_load_fields, + make_executable_tma_atom_or_pair, + pipeline_init_arrive_mixed_cga, + pipeline_init_wait_mixed_cga, + select_tma_atom_for_cluster, +) + + +__all__ = [ + "PipelineTmaUmmaMixedCga", + "TmaAtomOrPair", + "bind_executable_tma_load_fields", + "make_executable_tma_atom_or_pair", + "pipeline_init_arrive_mixed_cga", + "pipeline_init_wait_mixed_cga", + "select_tma_atom_for_cluster", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/__init__.py index 52a7a9daf028..51f74b497db7 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/__init__.py @@ -1,2 +1,25 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Rubin MegaMoE inference kernel components.""" + +from . import dynamic_mainloop +from .block_scaled_swap_ab_fc12_epilogue import GatedActEpilogueArgs, SwapABGatedActEpilogue +from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension, TensorRole +from .block_scaled_swap_ab_fc12_mainloop import BlockScaledSwapAbFc12Mainloop +from .block_scaled_swap_ab_fc12_mainloop_gen_specialized import BlockScaledSwapAbFc12MainloopGenSpecialized +from .block_scaled_swap_ab_mega_moe_kernel import BlockScaledSwapAbMegaMoeKernel +from .block_scaled_swap_ab_mega_moe_kernel_gen_specialized import BlockScaledSwapAbGenphaseMoeKernel +from .token_comm_gen_specialized import GenphaseTokenComm, GenphaseTokenCommArgs + + +__all__ = [ + "BlockScaledSwapAbFc12Extension", + "BlockScaledSwapAbFc12Mainloop", + "BlockScaledSwapAbFc12MainloopGenSpecialized", + "BlockScaledSwapAbGenphaseMoeKernel", + "BlockScaledSwapAbMegaMoeKernel", + "GatedActEpilogueArgs", + "GenphaseTokenComm", + "GenphaseTokenCommArgs", + "SwapABGatedActEpilogue", + "TensorRole", + "dynamic_mainloop", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/__init__.py index 52a7a9daf028..70999789b802 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/__init__.py @@ -1,2 +1,31 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Scheduler protocols and implementations.""" + +from .base import SchedulerBase, SchedulerConsumer, SchedulerWorkTileBase, WorkIdAcquisitionMode +from .clc_mixed_cga import ClcMixedCgaSchedulerWorker +from .fc12_mapping import ( + BlockPhase, + Fc12WorkTileState, + NonSwapAbFc12WorkTileInfo, + SwapAbFc12WorkTileInfo, + peek_ready_bit, +) +from .fc12_scheduler import BlackwellFusedFc12Scheduler, PhaseInterleavedFc12Scheduler +from .non_clc_mixed_cga import NonClcMixedCgaConfig, NonClcMixedCgaSchedulerWorker + + +__all__ = [ + "BlackwellFusedFc12Scheduler", + "BlockPhase", + "ClcMixedCgaSchedulerWorker", + "Fc12WorkTileState", + "NonSwapAbFc12WorkTileInfo", + "NonClcMixedCgaConfig", + "NonClcMixedCgaSchedulerWorker", + "PhaseInterleavedFc12Scheduler", + "SchedulerBase", + "SchedulerConsumer", + "SchedulerWorkTileBase", + "SwapAbFc12WorkTileInfo", + "WorkIdAcquisitionMode", + "peek_ready_bit", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py index 1315d5179679..7fbb2be22b76 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py @@ -19,20 +19,17 @@ from .....helpers.smem_workspace import SmemRegion, SmemWorkspace from .....quant_def import QuantKind from ....function_mapping import FunctionMapping -from ....schedulers.base import SchedulerConsumer -from ....schedulers.fc12_mapping import BlockPhase, SwapAbFc12WorkTileInfo +from ....schedulers import BlockPhase, SchedulerConsumer, SwapAbFc12WorkTileInfo from ..mega.block_scaled_swap_ab_fc12_epilogue import ( GatedActEpilogueArgs, SwapABFc1Epilogue, + SwapABFc2Epilogue as _SharedSwapABFc2Epilogue, + SwapABGatedActEpilogue as _SharedSwapABGatedActEpilogue, make_fc2_redg_process_pipeline, make_fc2_stg_process_pipeline, make_fc2_ublk_process_pipeline, ) -from ..mega.block_scaled_swap_ab_fc12_epilogue import SwapABFc2Epilogue as _SharedSwapABFc2Epilogue -from ..mega.block_scaled_swap_ab_fc12_epilogue import ( - SwapABGatedActEpilogue as _SharedSwapABGatedActEpilogue, -) -from ..mega.block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension +from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension from .local_routing import RoutingMetadata @@ -92,12 +89,8 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.fc2_use_ublk = self.fc2_use_bulk self.reduce_topk_in_epilogue = self.reduce_topk_in_kernel if not 1 <= self.fc1_epi_flag_batch <= self._EpilogueWarpCnt: - raise ValueError( - f"Asynchronous FC1 flag batch size must be in [1, {self._EpilogueWarpCnt}]." - ) - self.cluster_tile_intermediate_downproj = ( - self._EpilogueFc1IntermediateDownTileSize * self.cluster_shape_mn[0] - ) + raise ValueError(f"Asynchronous FC1 flag batch size must be in [1, {self._EpilogueWarpCnt}].") + self.cluster_tile_intermediate_downproj = self._EpilogueFc1IntermediateDownTileSize * self.cluster_shape_mn[0] atom_thr_size = 2 if self.use_2cta_instrs else 1 self.cta_tile_m = self._EpilogueFc2HiddenTileSize @@ -115,15 +108,9 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: if tmem_plan.accumulator_stage_stride_columns != tmem_plan.accumulator_stage_columns: raise ValueError("Rubin Local MegaMoE does not support overlapping accumulator stages.") if tmem_plan.accumulator_pipeline_stages != tmem_plan.accumulator_stage_count: - raise ValueError( - "Rubin Local MegaMoE requires one pipeline stage per disjoint accumulator stage." - ) + raise ValueError("Rubin Local MegaMoE requires one pipeline stage per disjoint accumulator stage.") self.num_accumulator_tmem_cols = tmem_plan.accumulator_columns - self.accumulator_shape = ( - self.cta_tile_m, - self.cta_tile_n, - tmem_plan.accumulator_stage_count, - ) + self.accumulator_shape = (self.cta_tile_m, self.cta_tile_n, tmem_plan.accumulator_stage_count) self.accumulator_stride = (1 << 16, 1, tmem_plan.accumulator_stage_stride_columns) self.fc2_hidden_needs_predicate = not ( @@ -131,36 +118,24 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: and self.hidden_size % (self.cta_tile_m * self.cluster_shape_mn[0]) == 0 ) self.intermediate_downproj = ( - self.intermediate_gateup_size // 2 - if isinstance(self.intermediate_gateup_size, int) - else None + self.intermediate_gateup_size // 2 if isinstance(self.intermediate_gateup_size, int) else None ) self.subtile_cnt = self.cta_tile_n // self._EpilogueTokenTileSize self.fc1_staging_stage_bytes = ( - self._EpilogueTokenTileSize - * self._EpilogueFc1IntermediateDownTileSize - * self.fc1_output_dtype.width - // 8 + self._EpilogueTokenTileSize * self._EpilogueFc1IntermediateDownTileSize * self.fc1_output_dtype.width // 8 ) self.fc1_staging_bytes = self.subtile_cnt * self.fc1_staging_stage_bytes self.fc1_amax_token_chunks = self._EpilogueTokenTileSize // 32 self.fc1_amax_slot_count = ( - self.fc1_amax_token_chunks * self._EpilogueWarpCnt * 32 - if self.needs_pair_amax_exchange - else 0 + self.fc1_amax_token_chunks * self._EpilogueWarpCnt * 32 if self.needs_pair_amax_exchange else 0 ) if self.fc1_amax_slot_count * 4 > self.fc1_staging_stage_bytes: raise ValueError("The paired amax exchange does not fit in one FC1 staging stage.") requested_fc2_tma_stages = impl_desc.get("fc2_tma_stages") - if ( - requested_fc2_tma_stages is not None - and not 1 <= requested_fc2_tma_stages <= self.subtile_cnt - ): - raise ValueError( - f"fc2_tma_stages must be in [1, {self.subtile_cnt}], got {requested_fc2_tma_stages}." - ) + if requested_fc2_tma_stages is not None and not 1 <= requested_fc2_tma_stages <= self.subtile_cnt: + raise ValueError(f"fc2_tma_stages must be in [1, {self.subtile_cnt}], got {requested_fc2_tma_stages}.") if self.fc2_use_bulk: single_stage_region = self._make_fc2_single_stage_region() if single_stage_region.nbytes % 16 != 0: @@ -170,27 +145,16 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: else: bf16_baseline_stages = min(2, self.subtile_cnt) bf16_stage_bytes = ( - self._EpilogueTokenTileSize - * self._EpilogueFc2HiddenTileSize - * cutlass.BFloat16.width - // 8 - ) - available_staging_bytes = max( - self.fc1_staging_bytes, bf16_baseline_stages * bf16_stage_bytes + self._EpilogueTokenTileSize * self._EpilogueFc2HiddenTileSize * cutlass.BFloat16.width // 8 ) - self.fc2_tma_stages = min( - self.subtile_cnt, available_staging_bytes // single_stage_region.nbytes - ) - self.fc2_staging_spec: Optional[SmemRegion] = self._make_fc2_staging_region( - self.fc2_tma_stages - ) + available_staging_bytes = max(self.fc1_staging_bytes, bf16_baseline_stages * bf16_stage_bytes) + self.fc2_tma_stages = min(self.subtile_cnt, available_staging_bytes // single_stage_region.nbytes) + self.fc2_staging_spec: Optional[SmemRegion] = self._make_fc2_staging_region(self.fc2_tma_stages) else: self.fc2_tma_stages = 0 self.fc2_staging_spec = None - def prepare_tma_store_params( - self, fc1_output_template: cute.Tensor - ) -> Tuple[cute.CopyAtom, cute.Tensor]: + def prepare_tma_store_params(self, fc1_output_template: cute.Tensor) -> Tuple[cute.CopyAtom, cute.Tensor]: fc1_operation = cpasync.CopyBulkTensorTileS2GOp() fc1_smem_layout = self.fc1_staged_smem_layout(1, without_stage_mode=True) fc1_tile = (self._EpilogueTokenTileSize, self._EpilogueFc1IntermediateDownTileSize) @@ -295,9 +259,7 @@ def run( cute.arch.cp_async_bulk_wait_group(0) wait_only_named_barrier.arrive_and_wait() if cur_was_linear1: - flag_tracker = fc1_epi.signal_fc1_done( - prev_work_tile_info, work_tile_info, flag_tracker - ) + flag_tracker = fc1_epi.signal_fc1_done(prev_work_tile_info, work_tile_info, flag_tracker) else: flag_tracker = fc2_epi.advance_fc1_flag_phase(work_tile_info, flag_tracker) flag_tracker.fire() @@ -364,13 +326,10 @@ def advance_fc1_flag_phase(self, next_work_tile_info, flag_tracker): @cute.jit def _make_output_router(self, work_tile_info: SwapAbFc12WorkTileInfo) -> "Fc2OutputRouter": task_tile_data_row_start = ( - work_tile_info.cumulative_data_physical_row - + work_tile_info.tile_n_idx * cutlass.Int32(self.cta_tile_n) + work_tile_info.cumulative_data_physical_row + work_tile_info.tile_n_idx * cutlass.Int32(self.cta_tile_n) ) hidden_base_this_cta_tile = work_tile_info.tile_m_idx * cutlass.Int32(self.cta_tile_m) - valid_hidden_this_cta_tile = ( - cutlass.Int32(self.fc2_output.shape[2]) - hidden_base_this_cta_tile - ) + valid_hidden_this_cta_tile = cutlass.Int32(self.fc2_output.shape[2]) - hidden_base_this_cta_tile if valid_hidden_this_cta_tile < 0: valid_hidden_this_cta_tile = 0 if valid_hidden_this_cta_tile > self._EpilogueFc2HiddenTileSize: @@ -425,14 +384,9 @@ def prefetch(self) -> "Fc2OutputRouter": if token_valid and hidden_valid: valid[iter_idx] = cutlass.Int32(1) metadata = RoutingMetadata.load( - self.routing_metadata.iterator.toint() - + Int64(token_in_tile) * Int64(RoutingMetadata.nbytes) - ) - dst_topk = ( - cutlass.Int32(0) - if cutlass.const_expr(self.reduce_topk_in_epilogue) - else metadata.topk_slot + self.routing_metadata.iterator.toint() + Int64(token_in_tile) * Int64(RoutingMetadata.nbytes) ) + dst_topk = cutlass.Int32(0) if cutlass.const_expr(self.reduce_topk_in_epilogue) else metadata.topk_slot dst_hidden = hidden_in_tile + self.hidden_base_this_cta_tile dst_ptrs[iter_idx] = cute.domain_offset( (Int64(metadata.token), dst_topk, dst_hidden), self.output @@ -441,9 +395,7 @@ def prefetch(self) -> "Fc2OutputRouter": return dataclasses.replace(self, dst_ptrs=dst_ptrs, valid=valid) @cute.jit - def get_data_dst( - self, iter_idx: Union[int, cutlass.Int32] - ) -> Tuple[cute.Pointer, cutlass.Int32]: + def get_data_dst(self, iter_idx: Union[int, cutlass.Int32]) -> Tuple[cute.Pointer, cutlass.Int32]: ptr = cute.make_ptr( self.output.element_type, self.dst_ptrs[iter_idx], diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_extension.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_extension.py new file mode 100644 index 000000000000..0bc99ae344f1 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_extension.py @@ -0,0 +1,7 @@ +"""Local MegaMoE shim for the compatible Rubin MegaMoE FC12 extension.""" + +# <<>> +from ..mega.block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension, TensorRole + + +__all__ = ["BlockScaledSwapAbFc12Extension", "TensorRole"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_mainloop.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_mainloop.py index 243336bbe19a..345070033ba4 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_mainloop.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_mainloop.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Mainloop component for the Rubin block-scaled swap-AB FC12 kernel.""" from typing import ClassVar, Tuple @@ -29,7 +26,7 @@ from .....quant_def import QuantKind from ....schedulers.base import SchedulerConsumer from ....schedulers.fc12_mapping import BlockPhase -from ..mega.block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension +from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension from .local_routing import RoutingMetadata from .tma_gather import sm107_tma_gather4_load @@ -37,20 +34,14 @@ class BlockScaledSwapAbFc12Mainloop(KernelComponent): """Own all Rubin load, MMA, pipeline, SMEM, and TMEM mainloop state.""" - ab_pipeline_mbarriers_region: ClassVar[str] = ( - "rubin.swap_ab_fc12.mainloop.ab_pipeline_mbarriers" - ) + ab_pipeline_mbarriers_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.ab_pipeline_mbarriers" a_smem_tensor_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.a_smem_tensor" b_smem_tensor_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.b_smem_tensor" sfa_smem_tensor_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.sfa_smem_tensor" sfb_smem_tensor_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.sfb_smem_tensor" - acc_pipeline_mbarriers_region: ClassVar[str] = ( - "rubin.swap_ab_fc12.mainloop.acc_pipeline_mbarriers" - ) + acc_pipeline_mbarriers_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.acc_pipeline_mbarriers" tmem_holding_buffer_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.tmem_holding_buffer" - tmem_deallocation_mbarrier_region: ClassVar[str] = ( - "rubin.swap_ab_fc12.mainloop.tmem_deallocation_mbarrier" - ) + tmem_deallocation_mbarrier_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.tmem_deallocation_mbarrier" tmem_allocation_barrier_id: ClassVar[int] = 2 fc1_gather_sync_barrier_id: ClassVar[int] = 3 @@ -95,9 +86,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.cluster_shape_mn = impl_desc["cluster_shape_mn"] self.use_2cta_instrs = impl_desc["use_2cta_instrs"] self.mainloop_smem_budget_bytes = impl_desc["mainloop_smem_budget_bytes"] - self.num_accumulator_consumer_warps_per_cta = impl_desc[ - "num_accumulator_consumer_warps_per_cta" - ] + self.num_accumulator_consumer_warps_per_cta = impl_desc["num_accumulator_consumer_warps_per_cta"] self.architecture = "sm_107" self.mma_cta_count = 2 if self.use_2cta_instrs else 1 @@ -118,15 +107,9 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: if tmem_plan.accumulator_stage_stride_columns != tmem_plan.accumulator_stage_columns: raise ValueError("Rubin Local MegaMoE does not support overlapping accumulator stages.") if tmem_plan.accumulator_pipeline_stages != tmem_plan.accumulator_stage_count: - raise ValueError( - "Rubin Local MegaMoE requires one pipeline stage per disjoint accumulator stage." - ) + raise ValueError("Rubin Local MegaMoE requires one pipeline stage per disjoint accumulator stage.") self.num_accumulator_tmem_cols = tmem_plan.accumulator_columns - self.accumulator_shape = ( - self.cta_tile_m, - self.cta_tile_n, - tmem_plan.accumulator_stage_count, - ) + self.accumulator_shape = (self.cta_tile_m, self.cta_tile_n, tmem_plan.accumulator_stage_count) self.accumulator_stride = (1 << 16, 1, tmem_plan.accumulator_stage_stride_columns) self.mma_instruction = Tcgen05MmaInstruction( @@ -139,12 +122,8 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: sfb_type=self.sf_dtype, sf_vec_size=self.sf_vec_size, ) - self.a_smem_alloc_dtype = tcgen05_smem_alloc_type( - self.a_dtype, self.b_dtype, self.architecture - ) - self.b_smem_alloc_dtype = tcgen05_smem_alloc_type( - self.b_dtype, self.a_dtype, self.architecture - ) + self.a_smem_alloc_dtype = tcgen05_smem_alloc_type(self.a_dtype, self.b_dtype, self.architecture) + self.b_smem_alloc_dtype = tcgen05_smem_alloc_type(self.b_dtype, self.a_dtype, self.architecture) self.sfb_instruction_shape_mnk = ( self.mma_instruction_mnk[0] // self.mma_cta_count, round_up(self.mma_instruction_mnk[1], 128), @@ -160,11 +139,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: sfb_type=self.sf_dtype, sf_vec_size=self.sf_vec_size, ) - self.mma_tiler_sfb = ( - self.sfb_instruction_shape_mnk[0], - self.sfb_instruction_shape_mnk[1], - self.mma_tile_k, - ) + self.mma_tiler_sfb = (self.sfb_instruction_shape_mnk[0], self.sfb_instruction_shape_mnk[1], self.mma_tile_k) self.cluster_layout_shape_vmnk = ( (self.mma_cta_count,), self.cluster_shape_mn[0] // self.mma_cta_count, @@ -187,17 +162,13 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: ): raise ValueError("Local MegaMoE AB stage transaction byte accounting is inconsistent.") if self.num_tma_load_bytes % 2 != 0: - raise ValueError( - "Local MegaMoE total AB transaction bytes must split evenly across two producer arrivals." - ) + raise ValueError("Local MegaMoE total AB transaction bytes must split evenly across two producer arrivals.") self.ab_pipeline_tx_count_per_producer = self.num_tma_load_bytes // 2 self.num_gather_b_warps = 4 self.gather_b_rows_per_cta = self.cta_tile_n // self.mma_cta_count self.gather_b_rows_per_warp = self.gather_b_rows_per_cta // self.num_gather_b_warps if self.gather_b_rows_per_warp % 4 != 0: - raise ValueError( - "Local MegaMoE gather-B rows per warp must be divisible by gather4 width." - ) + raise ValueError("Local MegaMoE gather-B rows per warp must be divisible by gather4 width.") self.gather_b_calls_per_warp = self.gather_b_rows_per_warp // 4 self.gather_b_stage_bytes = self.cta_tile_n * self.mma_tile_k * int(self.b_dtype.width) // 8 if self.gather_b_stage_bytes != self.b_stage_tma_bytes: @@ -211,9 +182,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: fc1_b_side_stage_bytes = self.gather_b_stage_bytes + self.fc1_sfb_stage_bytes fc2_b_side_stage_bytes = self.b_stage_tma_bytes + self.sfb_stage_tma_bytes if fc1_b_side_stage_bytes != fc2_b_side_stage_bytes: - raise ValueError( - "Local MegaMoE FC1 and FC2 B-side transaction bytes must be identical." - ) + raise ValueError("Local MegaMoE FC1 and FC2 B-side transaction bytes must be identical.") if self.fc1_sfb_stage_bytes <= 0: raise ValueError("Local MegaMoE FC1 SFB transaction byte count must be positive.") self.fc1_gather_sync_barrier = pipeline.NamedBarrier( @@ -223,9 +192,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: def _validate_configuration(self) -> None: if self.mma_k_mode != "2x": - raise ValueError( - f"Rubin Local MegaMoE only supports mma_k_mode='2x', got {self.mma_k_mode!r}." - ) + raise ValueError(f"Rubin Local MegaMoE only supports mma_k_mode='2x', got {self.mma_k_mode!r}.") for field_name, dimensions, expected_rank in ( ("mma_instruction_mnk", self.mma_instruction_mnk, 3), ("mma_tiler_mnk", self.mma_tiler_mnk, 3), @@ -233,10 +200,7 @@ def _validate_configuration(self) -> None: ): if len(dimensions) != expected_rank: raise ValueError(f"{field_name} must contain {expected_rank} dimensions.") - if not all( - isinstance(dimension, int) and not isinstance(dimension, bool) - for dimension in dimensions - ): + if not all(isinstance(dimension, int) and not isinstance(dimension, bool) for dimension in dimensions): raise TypeError(f"{field_name} dimensions must be Python integers.") if any(dimension <= 0 for dimension in dimensions): raise ValueError(f"{field_name} dimensions must be positive.") @@ -253,9 +217,7 @@ def _validate_configuration(self) -> None: f"got {self.mma_instruction_mnk[0]}." ) if self.mma_instruction_mnk[:2] != self.mma_tiler_mnk[:2]: - raise NotImplementedError( - "Rubin Local MegaMoE does not implement M/N instruction repetition or B-reuse." - ) + raise NotImplementedError("Rubin Local MegaMoE does not implement M/N instruction repetition or B-reuse.") if self.mma_tiler_mnk[1] not in (64, 128, 256): raise ValueError("Rubin Local MegaMoE supports tile N in (64, 128, 256).") if self.mma_tiler_mnk[2] % self.instruction_k != 0: @@ -276,9 +238,7 @@ def _validate_configuration(self) -> None: def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: """Register the selected Rubin mainloop SMEM plan.""" - smem_workspace.register_mbarrier( - self.ab_pipeline_mbarriers_region, self.num_ab_pipeline_stages * 2 - ) + smem_workspace.register_mbarrier(self.ab_pipeline_mbarriers_region, self.num_ab_pipeline_stages * 2) for region_name, region in ( (self.a_smem_tensor_region, self.a_smem_region), (self.b_smem_tensor_region, self.b_smem_region), @@ -293,37 +253,23 @@ def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: swizzle=region.swizzle, byte_alignment=region.byte_alignment, ) - smem_workspace.register_mbarrier( - self.acc_pipeline_mbarriers_region, self.num_accumulator_pipeline_stages * 2 - ) + smem_workspace.register_mbarrier(self.acc_pipeline_mbarriers_region, self.num_accumulator_pipeline_stages * 2) smem_workspace.register_mbarrier(self.tmem_deallocation_mbarrier_region, 1) - smem_workspace.register_tensor( - self.tmem_holding_buffer_region, cutlass.Int32, (1,), byte_alignment=8 - ) + smem_workspace.register_tensor(self.tmem_holding_buffer_region, cutlass.Int32, (1,), byte_alignment=8) def _select_ab_stage_plan(self) -> None: stage_one_regions = make_smem_layouts( - self.mma_instruction, - self.mma_tiler_mnk, - 1, - (self.a_major_mode, self.b_major_mode), - self.architecture, + self.mma_instruction, self.mma_tiler_mnk, 1, (self.a_major_mode, self.b_major_mode), self.architecture ) independent_sfb_region = make_smem_layouts( - self.sfb_mma_instruction, - self.mma_tiler_sfb, - 1, - (self.a_major_mode, self.b_major_mode), - self.architecture, + self.sfb_mma_instruction, self.mma_tiler_sfb, 1, (self.a_major_mode, self.b_major_mode), self.architecture )[3] if ( independent_sfb_region.shape != stage_one_regions[3].shape or independent_sfb_region.stride != stage_one_regions[3].stride or independent_sfb_region.nbytes != stage_one_regions[3].nbytes ): - raise ValueError( - "The Rubin independent SFB plan does not match the primary MMA SFB operand." - ) + raise ValueError("The Rubin independent SFB plan does not match the primary MMA SFB operand.") def packed_bytes(region, dtype) -> int: return (region.cosize * int(dtype.width) + 7) // 8 @@ -347,31 +293,22 @@ def packed_bytes(region, dtype) -> int: mbarrier_bytes = int(cutlass.Int64.width) // 8 ab_stage_cost_bytes = self.ab_stage_payload_bytes + 2 * mbarrier_bytes plan_tail_bytes = (2 * self.num_accumulator_pipeline_stages + 2) * mbarrier_bytes - self.num_ab_pipeline_stages = ( - self.mainloop_smem_budget_bytes - plan_tail_bytes - ) // ab_stage_cost_bytes + self.num_ab_pipeline_stages = (self.mainloop_smem_budget_bytes - plan_tail_bytes) // ab_stage_cost_bytes if self.num_ab_pipeline_stages < 1: raise ValueError( f"One AB stage needs {ab_stage_cost_bytes + plan_tail_bytes} bytes, exceeding the " f"{self.mainloop_smem_budget_bytes}-byte mainloop budget." ) - (self.a_smem_region, self.b_smem_region, self.sfa_smem_region, self.sfb_smem_region) = ( - make_smem_layouts( - self.mma_instruction, - self.mma_tiler_mnk, - self.num_ab_pipeline_stages, - (self.a_major_mode, self.b_major_mode), - self.architecture, - ) + (self.a_smem_region, self.b_smem_region, self.sfa_smem_region, self.sfb_smem_region) = make_smem_layouts( + self.mma_instruction, + self.mma_tiler_mnk, + self.num_ab_pipeline_stages, + (self.a_major_mode, self.b_major_mode), + self.architecture, ) self.selected_smem_bytes = sum( region.nbytes - for region in ( - self.a_smem_region, - self.b_smem_region, - self.sfa_smem_region, - self.sfb_smem_region, - ) + for region in (self.a_smem_region, self.b_smem_region, self.sfa_smem_region, self.sfb_smem_region) ) def make_tiled_mma(self) -> cute.TiledMma: @@ -439,11 +376,8 @@ def materialize_codegen_members(self) -> None: ("SFA", self.sfa_smem_layout_staged, self.sfa_smem_region), ("SFB", self.sfb_smem_layout_staged, self.sfb_smem_region), ): - if ( - materialized_layout.shape != planned_region.shape - or not strides_equal_ignoring_singletons( - planned_region.shape, materialized_layout.stride, planned_region.stride - ) + if materialized_layout.shape != planned_region.shape or not strides_equal_ignoring_singletons( + planned_region.shape, materialized_layout.stride, planned_region.stride ): raise ValueError( f"{tensor_name} SMEM plan mismatch: Python " @@ -469,24 +403,12 @@ def prepare_tma_load_params( sfa_stage_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0)) sfb_stage_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0)) - a_operation = sm100_utils.cluster_shape_to_tma_atom_A( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - b_operation = sm100_utils.cluster_shape_to_tma_atom_B( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - sfa_operation = sm100_utils.cluster_shape_to_tma_atom_A( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - sfb_operation = sm100_utils.cluster_shape_to_tma_atom_SFB( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - a_internal_type = ( - self.a_smem_alloc_dtype if self.a_smem_alloc_dtype is not self.a_dtype else None - ) - b_internal_type = ( - self.b_smem_alloc_dtype if self.b_smem_alloc_dtype is not self.b_dtype else None - ) + a_operation = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, self.tiled_mma.thr_id) + b_operation = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, self.tiled_mma.thr_id) + sfa_operation = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, self.tiled_mma.thr_id) + sfb_operation = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, self.tiled_mma.thr_id) + a_internal_type = self.a_smem_alloc_dtype if self.a_smem_alloc_dtype is not self.a_dtype else None + b_internal_type = self.b_smem_alloc_dtype if self.b_smem_alloc_dtype is not self.b_dtype else None fc1_a_atom, fc1_a_tensor = cute.nvgpu.make_tiled_tma_atom_A( a_operation, @@ -502,9 +424,7 @@ def prepare_tma_load_params( # through TMA TILE_GATHER4 while landing in the native K_SW128 B layout. fc1_b_2d = fc1_b[(None, None, 0)] gather_b_base = cute.make_layout((1, self.mma_tile_k), stride=(self.mma_tile_k, 1)) - gather_b_smem_layout = cute.make_composed_layout( - cute.make_swizzle(3, 4, 3), 0, gather_b_base - ) + gather_b_smem_layout = cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, gather_b_base) fc1_b_atom, fc1_b_tensor = cpasync.make_tiled_tma_atom( cpasync.CopyBulkTensorTileG2SOp(), fc1_b_2d, @@ -615,17 +535,10 @@ def assign_device_members( ) self.fc1_k_tile_count = ceil_div(bound_hidden_size, self.mma_tile_k) self.fc2_k_tile_count = ceil_div(bound_intermediate_gateup_size // 2, self.mma_tile_k) - self.sfb_vmnk_coord = ( - 0, - cta_coord_in_cluster[0], - cta_coord_in_cluster[1], - cta_coord_in_cluster[2], - ) + self.sfb_vmnk_coord = (0, cta_coord_in_cluster[0], cta_coord_in_cluster[1], cta_coord_in_cluster[2]) @cute.jit - def create_ab_pipeline( - self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer - ) -> pipeline.PipelineAsync: + def create_ab_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> pipeline.PipelineAsync: """Create the shared two-producer TMA-to-UMMA pipeline. Each producer arrival contributes the same accounting share. The two @@ -662,9 +575,7 @@ def advance_pipeline_state( def create_acc_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer): """Create the UMMA-to-epilogue accumulator pipeline.""" producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) - consumer_thread_count = ( - self.num_accumulator_consumer_warps_per_cta * 32 * self.mma_cta_count - ) + consumer_thread_count = self.num_accumulator_consumer_warps_per_cta * 32 * self.mma_cta_count consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, consumer_thread_count) return pipeline.PipelineUmmaAsync.create( barrier_storage=smem_workspace.ptr(self.acc_pipeline_mbarriers_region, smem_base), @@ -676,9 +587,7 @@ def create_acc_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Poi ) @cute.jit - def create_tmem_allocator( - self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer, *, allocator_warp_id: int - ): + def create_tmem_allocator(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer, *, allocator_warp_id: int): """Bind the SM107 TMEM allocator to finalized SMEM regions.""" allocation_barrier = pipeline.NamedBarrier( barrier_id=self.tmem_allocation_barrier_id, @@ -689,9 +598,7 @@ def create_tmem_allocator( barrier_for_retrieve=allocation_barrier, allocator_warp_id=allocator_warp_id, is_two_cta=self.use_2cta_instrs, - two_cta_tmem_dealloc_mbar_ptr=smem_workspace.ptr( - self.tmem_deallocation_mbarrier_region, smem_base - ), + two_cta_tmem_dealloc_mbar_ptr=smem_workspace.ptr(self.tmem_deallocation_mbarrier_region, smem_base), arch=self.architecture, ) @@ -705,18 +612,12 @@ def _s2t_copy_and_partition(self, smem_tensor: cute.Tensor, tmem_tensor: cute.Te mn_mode = cute.get(compact_smem_tensor.layout, mode=[0, 0]) mn_mode = cute.append(mn_mode, cute.make_layout((4,), stride=(0,))) - broadcast_layout = cute.append( - cute.group_modes(mn_mode, 0), cute.get(compact_smem_tensor.layout, mode=[0, 1]) - ) + broadcast_layout = cute.append(cute.group_modes(mn_mode, 0), cute.get(compact_smem_tensor.layout, mode=[0, 1])) broadcast_layout = cute.append( cute.group_modes(broadcast_layout, 0), cute.get(compact_smem_tensor.layout, mode=[1]) ) - broadcast_layout = cute.append( - broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[2]) - ) - broadcast_layout = cute.append( - broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[3]) - ) + broadcast_layout = cute.append(broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[2])) + broadcast_layout = cute.append(broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[3])) broadcast_smem_tensor = cute.make_tensor(compact_smem_tensor.iterator, broadcast_layout) partitioned_smem = thread_copy.partition_S(broadcast_smem_tensor) @@ -745,17 +646,10 @@ def run_tma_a( multicast_mask = None if cutlass.const_expr(self.is_a_mcast or self.use_2cta_instrs): multicast_mask = tma_multicast_mask( - self.cluster_shape_mn, - None, - self.cta_coord_in_cluster, - None, - self.use_2cta_instrs, - "a", + self.cluster_shape_mn, None, self.cta_coord_in_cluster, None, self.use_2cta_instrs, "a" ) - a_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_vmnk, (0, 0, None, 0)).shape - ) + a_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_vmnk, (0, 0, None, 0)).shape) sfa_cta_layout = a_cta_layout thread_mma = self.tiled_mma.get_slice(self.mma_cta_index) @@ -765,16 +659,10 @@ def run_tma_a( if is_fc1: iket.range_push("tma_weight_fc1") k_tile_count = self.fc1_k_tile_count - real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor( - "a", fc1_tma_a_tensor, work_tile - ) - real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfa", fc1_tma_sfa_tensor, work_tile - ) + real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor("a", fc1_tma_a_tensor, work_tile) + real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor("sfa", fc1_tma_sfa_tensor, work_tile) - global_a = cute.local_tile( - real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) - ) + global_a = cute.local_tile(real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None)) global_sfa = cute.local_tile( real_sfa, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) ) @@ -829,16 +717,10 @@ def run_tma_a( else: iket.range_push("tma_weight_fc2") k_tile_count = self.fc2_k_tile_count - real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor( - "a", fc2_tma_a_tensor, work_tile - ) - real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfa", fc2_tma_sfa_tensor, work_tile - ) + real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor("a", fc2_tma_a_tensor, work_tile) + real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor("sfa", fc2_tma_sfa_tensor, work_tile) - global_a = cute.local_tile( - real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) - ) + global_a = cute.local_tile(real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None)) global_sfa = cute.local_tile( real_sfa, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) ) @@ -920,28 +802,14 @@ def run_tma_b( sfb_multicast_mask = None if cutlass.const_expr(self.is_b_mcast or self.use_2cta_instrs): b_multicast_mask = tma_multicast_mask( - self.cluster_shape_mn, - None, - self.cta_coord_in_cluster, - None, - self.use_2cta_instrs, - "b", + self.cluster_shape_mn, None, self.cta_coord_in_cluster, None, self.use_2cta_instrs, "b" ) sfb_multicast_mask = tma_multicast_mask( - self.cluster_shape_mn, - None, - self.cta_coord_in_cluster, - None, - self.use_2cta_instrs, - "sfb", + self.cluster_shape_mn, None, self.cta_coord_in_cluster, None, self.use_2cta_instrs, "sfb" ) - b_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_vmnk, (0, None, 0, 0)).shape - ) - sfb_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape - ) + b_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_vmnk, (0, None, 0, 0)).shape) + sfb_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape) thread_mma = self.tiled_mma.get_slice(self.mma_cta_index) thread_mma_sfb = self.tiled_mma_sfb.get_slice(self.mma_cta_index) @@ -954,9 +822,7 @@ def run_tma_b( kernel_extension.wait_for_input(work_tile) iket.range_pop() k_tile_count = self.fc1_k_tile_count - real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfb", fc1_tma_sfb_tensor, work_tile - ) + real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor("sfb", fc1_tma_sfb_tensor, work_tile) global_sfb = cute.local_tile( real_sfb, cute.slice_(self.mma_tiler_sfb, (0, None, None)), (None, None, None) @@ -1013,16 +879,10 @@ def run_tma_b( iket.range_pop() k_tile_count = self.fc2_k_tile_count - real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor( - "b", fc2_tma_b_tensor, work_tile - ) - real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfb", fc2_tma_sfb_tensor, work_tile - ) + real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor("b", fc2_tma_b_tensor, work_tile) + real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor("sfb", fc2_tma_sfb_tensor, work_tile) - global_b = cute.local_tile( - real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None) - ) + global_b = cute.local_tile(real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None)) global_sfb = cute.local_tile( real_sfb, cute.slice_(self.mma_tiler_sfb, (0, None, None)), (None, None, None) ) @@ -1118,7 +978,9 @@ def run_tma_gather_b1( self.use_2cta_instrs, "b", ) - is_gather_b_mcast_leader = (not self.is_b_mcast) or (self.main_vmnk_coord[1] == Int32(0)) + is_gather_b_mcast_leader = (not self.is_b_mcast) or ( + self.main_vmnk_coord[1] == Int32(0) + ) work_tile = sched_consumer.consume_work() while work_tile.is_valid_tile: @@ -1127,8 +989,7 @@ def run_tma_gather_b1( iket.range_push("tma_gather_b1_fc1") kernel_extension.wait_for_input(work_tile) tile_pool_base = ( - work_tile.cumulative_data_physical_row - + work_tile.tile_n_idx * Int32(self.cta_tile_n) + work_tile.cumulative_data_physical_row + work_tile.tile_n_idx * Int32(self.cta_tile_n) ) cta_row_base = self.mma_cta_index * Int32(self.gather_b_rows_per_cta) warp_row_base = gather_warp_idx * Int32(rows_per_warp) @@ -1151,8 +1012,8 @@ def run_tma_gather_b1( stage_base_elements = ( stage_index * self.gather_b_rows_per_cta * self.mma_tile_k ) - warp_base_elements = gather_warp_idx * Int32( - rows_per_warp * self.mma_tile_k + warp_base_elements = ( + gather_warp_idx * Int32(rows_per_warp * self.mma_tile_k) ) for gather_idx in cutlass.range_constexpr(self.gather_b_calls_per_warp): row_start = gather_idx * 4 @@ -1204,33 +1065,24 @@ def run_mma( tmem_allocator.wait_for_alloc() accumulator_pointer = tmem_allocator.retrieve_ptr(self.acc_dtype) - accumulator_layout = cute.make_layout( - self.accumulator_shape, stride=self.accumulator_stride - ) + accumulator_layout = cute.make_layout(self.accumulator_shape, stride=self.accumulator_stride) accumulator_base = cute.make_tensor(accumulator_pointer, accumulator_layout) sfa_stage_smem_layout = cute.slice_( - blockscaled_utils.make_smem_layout_sfa( - tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1 - ), + blockscaled_utils.make_smem_layout_sfa(tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1), (None, None, None, 0), ) sfb_stage_smem_layout = cute.slice_( - blockscaled_utils.make_smem_layout_sfb( - tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1 - ), + blockscaled_utils.make_smem_layout_sfb(tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1), (None, None, None, 0), ) - sfa_pointer = cute.recast_ptr( - accumulator_pointer + self.num_accumulator_tmem_cols, dtype=self.sf_dtype - ) + sfa_pointer = cute.recast_ptr(accumulator_pointer + self.num_accumulator_tmem_cols, dtype=self.sf_dtype) sfa_layout = blockscaled_utils.make_tmem_layout_sfa( tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, sfa_stage_smem_layout ) tmem_sfa = cute.make_tensor(sfa_pointer, sfa_layout) sfb_pointer = cute.recast_ptr( - accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, - dtype=self.sf_dtype, + accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, dtype=self.sf_dtype ) sfb_layout = blockscaled_utils.make_tmem_layout_sfb( tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, sfb_stage_smem_layout @@ -1263,10 +1115,7 @@ def run_mma( if cutlass.const_expr(self.cta_tile_n == 64): sfb_shift = (work_tile.tile_n_idx % Int32(2)) * Int32(2) shifted_sfb_pointer = cute.recast_ptr( - accumulator_pointer - + self.num_accumulator_tmem_cols - + self.num_sfa_tmem_cols - + sfb_shift, + accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols + sfb_shift, dtype=self.sf_dtype, ) mma_tmem_sfb = cute.make_tensor(shifted_sfb_pointer, sfb_layout) @@ -1281,12 +1130,8 @@ def run_mma( next_full_status = ab_pipeline.consumer_try_wait(next_ab_pipeline_state) s2t_stage_coord = (None, None, None, None, ab_pipeline_state.index) - cute.copy( - tiled_copy_sfa, partitioned_smem_sfa[s2t_stage_coord], partitioned_tmem_sfa - ) - cute.copy( - tiled_copy_sfb, partitioned_smem_sfb[s2t_stage_coord], partitioned_tmem_sfb - ) + cute.copy(tiled_copy_sfa, partitioned_smem_sfa[s2t_stage_coord], partitioned_tmem_sfa) + cute.copy(tiled_copy_sfb, partitioned_smem_sfb[s2t_stage_coord], partitioned_tmem_sfb) a_stage_frag = fragment_a[(None, None, None, ab_pipeline_state.index)] b_stage_frag = fragment_b[(None, None, None, ab_pipeline_state.index)] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_local_mega_moe_kernel.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_local_mega_moe_kernel.py index 616bc2f576b4..3e86b1ad34e2 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_local_mega_moe_kernel.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_local_mega_moe_kernel.py @@ -1,11 +1,11 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Local MegaMoE kernel composition for Rubin.""" +import hashlib + from typing import ClassVar, Optional, Tuple import cuda.bindings.driver as cuda + import cutlass import cutlass.cute as cute import cutlass.pipeline as pipeline @@ -15,31 +15,28 @@ from cutlass.cutlass_dsl import Int32 from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait -from .....api import ( - ImplDesc, - KernelClass, - OptionalRequirement, - ProblemDesc, - StaticOrRuntimeIntegerType, -) -from .....helpers.cute_py_helpers import ( - Tcgen05MmaInstruction, - make_tcgen05_tmem_plan, - tcgen05_block_scaled_acc_dtype, -) +from .....api import ImplDesc, KernelClass, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType +from .....helpers.cute_py_helpers import Tcgen05MmaInstruction, make_tcgen05_tmem_plan, tcgen05_block_scaled_acc_dtype from .....helpers.device_workspace import DeviceWorkspace +from .....helpers.dlb_expert_count import resolve_dlb_expert_counts from .....helpers.iket_compat import iket +from .....helpers.megamoe_aot import make_megamoe_aot_callable from .....helpers.smem_workspace import SmemWorkspace +from .....helpers.topk_reduce import TopkReduce from .....helpers.utils import ceil_div, round_up +from .....helpers.helper_weight_ready_gate import ( + derive_helper_weight_ready_topology, + make_helper_weight_ready_generation_gate, +) from .....quant_def import CombineFormat, QuantKind -from ....blackwell.inference.mega.topk_reduce import TopkReduce from ....schedulers.base import WorkIdAcquisitionMode from ....schedulers.fc12_scheduler import BlackwellFusedFc12Scheduler, PhaseInterleavedFc12Scheduler -from ..mega.block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension from .block_scaled_swap_ab_fc12_epilogue import GatedActEpilogueArgs, SwapABGatedActEpilogue +from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension from .block_scaled_swap_ab_fc12_mainloop import BlockScaledSwapAbFc12Mainloop from .local_routing import LocalRouting, LocalRoutingArgs + _aot_symbol_prefix = "rubin_local_mega_moe_aot" @@ -97,10 +94,16 @@ def impl_desc_require(cls) -> dict[str, type]: "launch_cluster_count": int, "routing_flag_batch": int, "reduce_topk_in_kernel": bool, + "helper_expert_count": OptionalRequirement(int), + "total_helper_slots": OptionalRequirement(int), + "weight_ready_source_count": OptionalRequirement(int), } def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self._validate_desc_inputs(problem_desc, impl_desc) + problem_desc, impl_desc = resolve_dlb_expert_counts( + problem_desc, impl_desc, world_size=1 + ) self.expert_count = problem_desc["expert_count"] self.intermediate_gateup_size = problem_desc["intermediate_gateup_size"] @@ -141,6 +144,9 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.launch_cluster_count = impl_desc["launch_cluster_count"] self.routing_flag_batch = impl_desc["routing_flag_batch"] self.reduce_topk_in_kernel = impl_desc["reduce_topk_in_kernel"] + self.helper_expert_count = impl_desc.get("helper_expert_count", 0) + self.weight_ready_source_count = impl_desc.get("weight_ready_source_count", 1) + self._validate_helper_configuration() self.occupancy = 1 self.architecture = "sm_107" @@ -170,9 +176,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: "promised_launchable_sm_count": ( self.launch_cluster_count * self.cluster_shape_mn[0] * self.cluster_shape_mn[1] ), - "router_smem_limit_bytes": cutlass.memory.get_smem_capacity_in_bytes( - self.architecture - ), + "router_smem_limit_bytes": utils.get_smem_capacity_in_bytes(self.architecture), } ) resolved_problem_desc = ProblemDesc( @@ -185,9 +189,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.local_routing = LocalRouting(resolved_problem_desc, local_routing_impl_desc) self.max_routed_tokens = self.local_routing.worst_case_token_count - fc12_problem_desc = ProblemDesc( - {**resolved_problem_desc, "expert_count": self.expert_count} - ) + fc12_problem_desc = ProblemDesc({**resolved_problem_desc, "expert_count": self.expert_count}) resolved_impl_desc = ImplDesc( { **impl_desc, @@ -201,9 +203,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: } ) scheduler_type = ( - BlackwellFusedFc12Scheduler - if self.schedule_mode == "grouped" - else PhaseInterleavedFc12Scheduler + BlackwellFusedFc12Scheduler if self.schedule_mode == "grouped" else PhaseInterleavedFc12Scheduler ) self.scheduler = scheduler_type(fc12_problem_desc, resolved_impl_desc) self.epilogue = SwapABGatedActEpilogue(fc12_problem_desc, resolved_impl_desc) @@ -214,15 +214,19 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: if self.epilogue.sf_vec_size != self.sf_vec_size: raise ValueError("Epilogue and Mainloop scale vector sizes must match.") self._device_workspace = self._build_device_workspace() - self._mainloop, self._smem_workspace = self._build_mainloop_and_smem( - fc12_problem_desc, resolved_impl_desc - ) + self._mainloop, self._smem_workspace = self._build_mainloop_and_smem(fc12_problem_desc, resolved_impl_desc) self._topk_reduce = ( - None - if self.reduce_topk_in_kernel - else TopkReduce(self.hidden_size, self.topk, self.combine_format) + None if self.reduce_topk_in_kernel else TopkReduce(self.hidden_size, self.topk, self.combine_format) ) + def _validate_helper_configuration(self) -> None: + if type(self.helper_expert_count) is not int or self.helper_expert_count < 0: + raise ValueError("helper_expert_count must be a non-negative exact int") + if type(self.weight_ready_source_count) is not int or self.weight_ready_source_count <= 0: + raise ValueError("weight_ready_source_count must be a positive exact int") + if self.helper_expert_count > 0: + derive_helper_weight_ready_topology(self.expert_count, self.helper_expert_count) + def _validate_geometry(self) -> None: static_expert_dimensions = ( isinstance(self.expert_count, int), @@ -230,21 +234,15 @@ def _validate_geometry(self) -> None: isinstance(self.hidden_size, int), ) if any(static_expert_dimensions) and not all(static_expert_dimensions): - raise ValueError( - "Local MegaMoE expert dimensions must be either all static or all runtime." - ) + raise ValueError("Local MegaMoE expert dimensions must be either all static or all runtime.") if not all(static_expert_dimensions): - raise NotImplementedError( - "The Local MegaMoE kernel currently requires static expert dimensions." - ) + raise NotImplementedError("The Local MegaMoE kernel currently requires static expert dimensions.") if self.expert_count <= 0 or self.intermediate_gateup_size <= 0 or self.hidden_size <= 0: raise ValueError("Local MegaMoE expert dimensions must be positive.") if self.topk <= 0 or self.topk > self.expert_count: raise ValueError("topk must be positive and no greater than expert_count.") if self.topk_index_dtype not in (cutlass.Int32, cutlass.Int64): - raise ValueError( - f"topk_index_dtype must be Int32 or Int64, got {self.topk_index_dtype}." - ) + raise ValueError(f"topk_index_dtype must be Int32 or Int64, got {self.topk_index_dtype}.") if self.reduce_topk_in_kernel and not self.apply_topk_at_fc1: raise ValueError("In-kernel top-k reduction requires apply_topk_at_fc1=True.") if self.intermediate_gateup_size % 2 != 0: @@ -261,12 +259,8 @@ def _validate_geometry(self) -> None: raise ValueError( f"schedule_policy mode must be 'grouped' or 'phase_interleave', got {self.schedule_mode!r}." ) - if self.hint is not None and ( - isinstance(self.hint, bool) or not isinstance(self.hint, int) or self.hint <= 0 - ): - raise ValueError( - f"schedule_policy hint must be a positive Python int, got {self.hint!r}." - ) + if self.hint is not None and (isinstance(self.hint, bool) or not isinstance(self.hint, int) or self.hint <= 0): + raise ValueError(f"schedule_policy hint must be a positive Python int, got {self.hint!r}.") if self.schedule_mode == "phase_interleave" and self.work_id_mode != "atomic_counter": raise ValueError("phase_interleave currently requires work_id_mode='atomic_counter'.") if not 1 <= self.fc1_epi_flag_batch <= len(self.epilogue_warp_ids): @@ -282,17 +276,12 @@ def _validate_geometry(self) -> None: ): if len(dimensions) != expected_rank: raise ValueError(f"{field_name} must contain {expected_rank} dimensions.") - if not all( - isinstance(dimension, int) and not isinstance(dimension, bool) - for dimension in dimensions - ): + if not all(isinstance(dimension, int) and not isinstance(dimension, bool) for dimension in dimensions): raise TypeError(f"{field_name} dimensions must be Python integers.") if any(dimension <= 0 for dimension in dimensions): raise ValueError(f"{field_name} dimensions must be positive.") if self.mma_k_mode != "2x": - raise ValueError( - f"Rubin Local MegaMoE requires mma_k_mode='2x', got {self.mma_k_mode!r}." - ) + raise ValueError(f"Rubin Local MegaMoE requires mma_k_mode='2x', got {self.mma_k_mode!r}.") mma_m, mma_n, mma_k = self.mma_tiler_mnk instruction_m, instruction_n, instruction_k = self.mma_instruction_mnk @@ -309,9 +298,7 @@ def _validate_geometry(self) -> None: f"{self.quant_kind} requires Rubin 2x instruction K={expected_instruction_k}, got {instruction_k}." ) if (instruction_m, instruction_n) != (mma_m, mma_n): - raise NotImplementedError( - "Rubin Local MegaMoE does not implement M/N instruction repetition or B-reuse." - ) + raise NotImplementedError("Rubin Local MegaMoE does not implement M/N instruction repetition or B-reuse.") if mma_n not in (64, 128, 256): raise ValueError(f"mma_tiler N must be 64, 128, or 256, got {mma_n}.") if mma_k % instruction_k != 0: @@ -323,9 +310,7 @@ def _validate_geometry(self) -> None: "Rubin Local MegaMoE excludes NVFP4 N256 K512 because its multi-window path has no measured gain." ) if self.intermediate_gateup_size % (self.sf_vec_size * 4) != 0: - raise ValueError( - "The intermediate dimension must be divisible by four scale-factor vectors." - ) + raise ValueError("The intermediate dimension must be divisible by four scale-factor vectors.") if cluster_n != 1: raise ValueError(f"The swap-AB FC12 path requires cluster N=1, got {cluster_n}.") if cluster_m <= 0 or cluster_m > 4 or cluster_m & (cluster_m - 1): @@ -342,9 +327,7 @@ def _resolve_schedule_hint(self) -> None: blocks_fc1 = ceil_div(self.intermediate_gateup_size, cluster_feature_tile) blocks_fc2 = ceil_div(self.hidden_size, cluster_feature_tile) # Cover the worst token-block alignment of one full persistent-cluster FC2 claim wave. - max_dependent_token_blocks = ceil_div( - self.launch_cluster_count + blocks_fc2 - 1, blocks_fc2 - ) + max_dependent_token_blocks = ceil_div(self.launch_cluster_count + blocks_fc2 - 1, blocks_fc2) required_fc1_work = max_dependent_token_blocks * blocks_fc1 raw_minimum_hint = max(1, ceil_div(required_fc1_work, self.launch_cluster_count)) minimum_safe_hint = round_up(raw_minimum_hint, self.fc1_epi_flag_batch) @@ -419,6 +402,7 @@ def dtype_name(dtype: type) -> str: f"{dtype_name(self.a_dtype)}_{dtype_name(self.b_dtype)}_{dtype_name(self.acc_dtype)}_" f"sfvec{self.sf_vec_size}_a{self.a_major_mode.name.lower()}_b{self.b_major_mode.name.lower()}_" f"e{self.expert_count}_topk{self.topk}_" + f"helpers{self.helper_expert_count}_weightsources{self.weight_ready_source_count}_" f"topkidx{dtype_name(self.topk_index_dtype)}_" f"h{self.hidden_size}_i{self.intermediate_gateup_size}_maxtoken{self.max_tokens}_" f"inst{instruction}_{self.mma_k_mode}_tile{tile}_" @@ -435,6 +419,7 @@ def dtype_name(dtype: type) -> str: f"situ{self.situ_beta}x{self.situ_linear_beta}_" f"{'apply_topk_fc1' if self.apply_topk_at_fc1 else 'apply_topk_fc2'}_" f"{'inkernel_redg' if self.reduce_topk_in_kernel else 'separate_reduce'}" + + ("_dlb_v1" if self.helper_expert_count > 0 else "") ) def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): @@ -453,14 +438,10 @@ def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): extents = tuple( - sym_int64(divisibility=math.gcd(int(extent), 128)) - if axis in dynamic_axes - else int(extent) + sym_int64(divisibility=math.gcd(int(extent), 128)) if axis in dynamic_axes else int(extent) for axis, extent in enumerate(shape) ) - return make_fake_compact_tensor( - dtype, extents, stride_order=stride_order, assumed_align=alignment - ) + return make_fake_compact_tensor(dtype, extents, stride_order=stride_order, assumed_align=alignment) tokens = self.max_tokens hidden = self.hidden_size @@ -469,20 +450,14 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): experts = self.expert_count sf_vec_size = self.sf_vec_size # Weight SF is atom-swizzled and opaque, so its extent is whatever the caller staged, not a kernel choice. - fc1_weight_sf_columns = round_up(intermediate_gateup, 128) * round_up( - hidden // sf_vec_size, 4 - ) - fc2_weight_sf_columns = round_up(hidden, 128) * round_up( - intermediate_downproj // sf_vec_size, 4 - ) + fc1_weight_sf_columns = round_up(intermediate_gateup, 128) * round_up(hidden // sf_vec_size, 4) + fc2_weight_sf_columns = round_up(hidden, 128) * round_up(intermediate_downproj // sf_vec_size, 4) # TopkReduce and the in-kernel REDG both take their output dtype from this tensor. output_dtype = cutlass.BFloat16 fake_arguments = dict( # Swap-AB puts the weights on the MMA A operand, so activation uses the routing input dtype. - activation=fake_tensor( - self.local_routing.activation_dtype, (tokens, hidden), (1, 0), {0}, 16 - ), + activation=fake_tensor(self.local_routing.activation_dtype, (tokens, hidden), (1, 0), {0}, 16), activation_sf=fake_tensor( self.local_routing.activation_sf_dtype, (tokens, self.local_routing.activation_sf_hidden_padded), @@ -492,18 +467,10 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): ), topk_indices=fake_tensor(self.topk_index_dtype, (tokens, self.topk), (1, 0), {0}, 16), topk_scores=fake_tensor(cutlass.Float32, (tokens, self.topk), (1, 0), {0}, 4), - fc1_weight=fake_tensor( - self.a_dtype, (experts, hidden, intermediate_gateup), (2, 0, 1), {0, 2}, 16 - ), - fc1_weight_sf=fake_tensor( - self.sf_dtype, (experts, fc1_weight_sf_columns), (1, 0), {0}, 16 - ), - fc2_weight=fake_tensor( - self.a_dtype, (experts, intermediate_downproj, hidden), (2, 0, 1), {0, 2}, 16 - ), - fc2_weight_sf=fake_tensor( - self.sf_dtype, (experts, fc2_weight_sf_columns), (1, 0), {0}, 16 - ), + fc1_weight=fake_tensor(self.a_dtype, (experts, hidden, intermediate_gateup), (2, 0, 1), {0, 2}, 16), + fc1_weight_sf=fake_tensor(self.sf_dtype, (experts, fc1_weight_sf_columns), (1, 0), {0}, 16), + fc2_weight=fake_tensor(self.a_dtype, (experts, intermediate_downproj, hidden), (2, 0, 1), {0, 2}, 16), + fc2_weight_sf=fake_tensor(self.sf_dtype, (experts, fc2_weight_sf_columns), (1, 0), {0}, 16), output_activation=fake_tensor(output_dtype, (tokens, hidden), (1, 0), {0}, 16), # Opaque byte workspaces: a bare base pointer, like to_cute_ptr. The 128-byte promise matches the # alignment the registered regions assume, so the caller must hand over 128-byte-aligned bases. @@ -517,24 +484,34 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): fc1_norm_const=fake_tensor(cutlass.Float32, (experts,), (0,), set(), 4), ) + if self.helper_expert_count > 0: + fake_arguments.update( + hot_expert_weight_ready_flags=make_ptr( + cutlass.Uint64, 0, AddressSpace.gmem, assumed_align=8 + ), + hot_expert_weight_ready_generation=cutlass.Uint64(1), + ) + compiled = cute.compile[cute.EnableTVMFFI(True)](self, **fake_arguments) if out_path is None: return compiled - compiled.export_to_c( - out_path, function_name=_aot_symbol_prefix, export_only_tvm_ffi_symbols=True - ) + compiled.export_to_c(out_path, function_name=self._aot_symbol_name(), export_only_tvm_ffi_symbols=True) return out_path - @staticmethod - def load_compiled(path: str): + def _aot_symbol_name(self) -> str: + identity = hashlib.sha256(self.name().encode()).hexdigest()[:16] + return f"{_aot_symbol_prefix}_{identity}" + + def load_compiled(self, path: str): from cutlass.cute.runtime import load_module - return load_module(path, enable_tvm_ffi=True)[_aot_symbol_prefix] + function = getattr(load_module(path, enable_tvm_ffi=True), self._aot_symbol_name()) + return make_megamoe_aot_callable(function, self) def _build_mainloop_and_smem( self, problem_desc: ProblemDesc, resolved_impl_desc: ImplDesc ) -> Tuple[BlockScaledSwapAbFc12Mainloop, SmemWorkspace]: - smem_limit = cutlass.memory.get_smem_capacity_in_bytes(self.architecture) // self.occupancy + smem_limit = utils.get_smem_capacity_in_bytes(self.architecture) // self.occupancy smem_workspace = SmemWorkspace() self.local_routing.register_smem_regions(smem_workspace) self.scheduler.register_smem_regions(smem_workspace) @@ -545,9 +522,7 @@ def _build_mainloop_and_smem( if mainloop_smem_budget_bytes <= 0: raise ValueError("Fixed kernel components leave no SMEM budget for the Mainloop.") - mainloop_impl_desc = ImplDesc( - {**resolved_impl_desc, "mainloop_smem_budget_bytes": mainloop_smem_budget_bytes} - ) + mainloop_impl_desc = ImplDesc({**resolved_impl_desc, "mainloop_smem_budget_bytes": mainloop_smem_budget_bytes}) mainloop = BlockScaledSwapAbFc12Mainloop(problem_desc, mainloop_impl_desc) mainloop.register_smem_regions(smem_workspace) @@ -572,16 +547,28 @@ def __call__( fc1_alpha: Optional[cute.Tensor] = None, # (local_experts,) fc2_alpha: Optional[cute.Tensor] = None, # (local_experts,) fc1_norm_const: Optional[cute.Tensor] = None, # (local_experts,) + hot_expert_weight_ready_flags: Optional[cute.Pointer] = None, # uint64[weight_ready_source_count] + hot_expert_weight_ready_generation: Optional[cutlass.Uint64] = None, ) -> None: """Launch router, fused MegaMoE compute, and optional TopK reduce.""" + helper_ready_bundle = ( + hot_expert_weight_ready_flags, + hot_expert_weight_ready_generation, + ) + if cutlass.const_expr( + (self.helper_expert_count > 0 and not all(value is not None for value in helper_ready_bundle)) + or (self.helper_expert_count == 0 and any(value is not None for value in helper_ready_bundle)) + ): + raise ValueError( + "helper READY flags/generation must both be present iff helper_expert_count > 0" + ) + def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: return cute.make_tensor(tensor.iterator, cute.make_layout(shape, stride=tensor.stride)) if cutlass.const_expr(topk_indices.element_type is not self.topk_index_dtype): - raise TypeError( - f"topk_indices must be the declared {self.topk_index_dtype}, got {topk_indices.dtype}." - ) + raise TypeError(f"topk_indices must be the declared {self.topk_index_dtype}, got {topk_indices.dtype}.") for operand_name, operand, expected_dtype in ( ("activation", activation, self.b_dtype), ("activation_sf", activation_sf, self.sf_dtype), @@ -654,19 +641,14 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: padded_intermediate_gateup = cute.round_up(intermediate_gateup, sf_vec_size * 4) expected_fc1_weight_sf_columns = padded_intermediate_gateup * padded_hidden // sf_vec_size if cutlass.const_expr( - isinstance(fc1_weight_sf.shape[1], int) - and isinstance(expected_fc1_weight_sf_columns, int) + isinstance(fc1_weight_sf.shape[1], int) and isinstance(expected_fc1_weight_sf_columns, int) ): if cutlass.const_expr(fc1_weight_sf.shape[1] != expected_fc1_weight_sf_columns): raise ValueError("fc1_weight_sf has an incompatible column count.") fc1_sfa = cute.make_tensor( fc1_weight_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - cutlass.Int32(padded_intermediate_gateup), - cutlass.Int32(padded_hidden), - cutlass.Int32(experts), - ), + (cutlass.Int32(padded_intermediate_gateup), cutlass.Int32(padded_hidden), cutlass.Int32(experts)), sf_vec_size, ), ) @@ -684,31 +666,21 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: padded_fc2_intermediate = cute.round_up(intermediate_fc2, sf_vec_size * 4) expected_fc2_weight_sf_columns = padded_fc2_hidden * padded_fc2_intermediate // sf_vec_size if cutlass.const_expr( - isinstance(fc2_weight_sf.shape[1], int) - and isinstance(expected_fc2_weight_sf_columns, int) + isinstance(fc2_weight_sf.shape[1], int) and isinstance(expected_fc2_weight_sf_columns, int) ): if cutlass.const_expr(fc2_weight_sf.shape[1] != expected_fc2_weight_sf_columns): raise ValueError("fc2_weight_sf has an incompatible column count.") fc2_sfa = cute.make_tensor( fc2_weight_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - cutlass.Int32(padded_fc2_hidden), - cutlass.Int32(padded_fc2_intermediate), - cutlass.Int32(experts_fc2), - ), + (cutlass.Int32(padded_fc2_hidden), cutlass.Int32(padded_fc2_intermediate), cutlass.Int32(experts_fc2)), sf_vec_size, ), ) fc2_sfb = cute.make_tensor( fc1_output_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - fc1_output_sf.shape[0], - cutlass.Int32(fc1_output_sf.shape[1] * sf_vec_size), - singleton, - ), - sf_vec_size, + (fc1_output_sf.shape[0], cutlass.Int32(fc1_output_sf.shape[1] * sf_vec_size), singleton), sf_vec_size ), ) if cutlass.const_expr(self.reduce_topk_in_kernel): @@ -716,21 +688,14 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: output_activation.iterator, cute.make_layout( (self.max_tokens, 1, hidden), - stride=( - output_activation.stride[0], - output_activation.stride[0], - output_activation.stride[1], - ), + stride=(output_activation.stride[0], output_activation.stride[0], output_activation.stride[1]), ), ) else: - pre_reduced_activation = self.local_routing.pre_reduced_activation_tensor( - self._device_workspace - ) + pre_reduced_activation = self.local_routing.pre_reduced_activation_tensor(self._device_workspace) fc2_output = rewrite_tensor_shape( - pre_reduced_activation, - (pre_reduced_activation.shape[0], pre_reduced_activation.shape[1], hidden), + pre_reduced_activation, (pre_reduced_activation.shape[0], pre_reduced_activation.shape[1], hidden) ) self._mainloop.materialize_codegen_members() @@ -762,9 +727,7 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: fc2_sfb=fc2_sfb, ) - fc1_output_tma_atom, fc1_output_tma_tensor = self.epilogue.prepare_tma_store_params( - fc1_output_gemm - ) + fc1_output_tma_atom, fc1_output_tma_tensor = self.epilogue.prepare_tma_store_params(fc1_output_gemm) grid = self.scheduler.get_grid_shape(max_active_clusters=self.launch_cluster_count) self._device_workspace.remove_device_members() @@ -797,6 +760,8 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: fc1_alpha, fc2_alpha, fc1_norm_const, + hot_expert_weight_ready_flags, + hot_expert_weight_ready_generation, ).launch( grid=grid, block=[self.threads_per_cta, 1, 1], @@ -807,9 +772,7 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: ) if cutlass.const_expr(not self.reduce_topk_in_kernel): reduce_scores = None if cutlass.const_expr(self.apply_topk_at_fc1) else topk_scores - self._topk_reduce( - pre_reduced_activation, None, output_activation, reduce_scores, stream - ) + self._topk_reduce(pre_reduced_activation, None, output_activation, reduce_scores, stream) @cute.kernel def _kernel( @@ -842,6 +805,8 @@ def _kernel( fc1_alpha: Optional[cute.Tensor], fc2_alpha: Optional[cute.Tensor], fc1_norm_const: Optional[cute.Tensor], + hot_expert_weight_ready_flags: Optional[cute.Pointer], + hot_expert_weight_ready_generation: Optional[cutlass.Uint64], ): """Compose local routing, scheduler, mainloop, and epilogue.""" self._mainloop.materialize_codegen_members() @@ -863,8 +828,7 @@ def _kernel( fc1_output_sf = cute.make_tensor( fc1_output_sf_storage.iterator, blockscaled_utils.tile_atom_to_shape_SF( - (fc1_output_sf_storage.shape[0], intermediate_gateup // 2, cutlass.Int32(1)), - self.epilogue.sf_vec_size, + (fc1_output_sf_storage.shape[0], intermediate_gateup // 2, cutlass.Int32(1)), self.epilogue.sf_vec_size ), ) @@ -918,22 +882,24 @@ def _kernel( ) scheduler = self.scheduler - fc2_spin_threshold = ( - intermediate_gateup + self._mainloop.cta_tile_m - 1 - ) // self._mainloop.cta_tile_m + fc2_spin_threshold = (intermediate_gateup + self._mainloop.cta_tile_m - 1) // self._mainloop.cta_tile_m + helper_weight_ready_gate = None + if cutlass.const_expr(self.helper_expert_count > 0): + helper_weight_ready_gate = make_helper_weight_ready_generation_gate( + memory_slot_count=self.expert_count, + helper_count=self.helper_expert_count, + source_count=self.weight_ready_source_count, + terminal_flags=hot_expert_weight_ready_flags, + expected_generation=hot_expert_weight_ready_generation, + ) kernel_extension = BlockScaledSwapAbFc12Extension( sf_vec_size=self.sf_vec_size, fc1_done_counter_pointer=fc1_done_counter.iterator, fc2_spin_threshold=fc2_spin_threshold, - fc1_ready_counter_pointer=self.local_routing.fc1_ready_counter_pointer( - self._device_workspace - ), + fc1_ready_counter_pointer=self.local_routing.fc1_ready_counter_pointer(self._device_workspace), ) optional_epilogue_args = GatedActEpilogueArgs( - fc1_alpha=fc1_alpha, - fc2_alpha=fc2_alpha, - fc1_norm_const=fc1_norm_const, - topk_scores=pool_topk_scores, + fc1_alpha=fc1_alpha, fc2_alpha=fc2_alpha, fc1_norm_const=fc1_norm_const, topk_scores=pool_topk_scores ) pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) @@ -951,9 +917,15 @@ def _kernel( iket.range_push("scheduler.gen_work") work_tile = scheduler.gen_next_work() iket.range_pop() + helper_ready_latched = Int32(0) while work_tile.is_valid_tile: iket.range_push("scheduler.publish_work") - scheduler.publish_work(kernel_extension.prepare_work_tile(work_tile)) + prepared_tile = kernel_extension.prepare_work_tile(work_tile) + if cutlass.const_expr(helper_weight_ready_gate is not None): + helper_ready_latched = helper_weight_ready_gate.wait_before_first_helper_publish( + prepared_tile.expert_idx, helper_ready_latched + ) + scheduler.publish_work(prepared_tile) iket.range_pop() iket.range_push("scheduler.gen_work") work_tile = scheduler.gen_next_work() @@ -1043,9 +1015,7 @@ def _kernel( optional_epilogue_args, ) tmem_allocator.relinquish_alloc_permit() - tmem_allocator.free( - tmem_allocator.retrieve_ptr(self.acc_dtype), self._mainloop.num_tmem_alloc_cols - ) + tmem_allocator.free(tmem_allocator.retrieve_ptr(self.acc_dtype), self._mainloop.num_tmem_alloc_cols) iket.range_pop() if warp_idx >= Int32(self.transfer_warp_idx_start): @@ -1074,9 +1044,7 @@ def _kernel( iket.range_push("local_mega.kernel_tail_wait") cute.arch.sync_threads() iket.range_pop() - if (warp_idx >= Int32(self.transfer_warp_idx_start)) & ( - warp_idx < Int32(self.token_in_end_warp_idx) - ): + if (warp_idx >= Int32(self.transfer_warp_idx_start)) & (warp_idx < Int32(self.token_in_end_warp_idx)): iket.range_push("local_mega.tail_reset") self.local_routing.reset_tail() iket.range_pop() diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/local_routing.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/local_routing.py index 3ed254dcfc0c..fd3ceb087e89 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/local_routing.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/local_routing.py @@ -1,12 +1,10 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Routing support for the Rubin Local MegaMoE kernel.""" import dataclasses from typing import Callable, ClassVar, Optional, Tuple, Union import cuda.bindings.driver as cuda + import cutlass import cutlass.cute as cute import cutlass.pipeline as pipeline @@ -20,7 +18,13 @@ from .....helpers.dsl_helpers import mark_alignment from .....helpers.flag_batch import make_flag_batch_tracker from .....helpers.iket_compat import iket -from .....helpers.ptx_helpers import cp_async_bulk_s2g, nanosleep, stg_b64, stg_f32, tma_load_1d +from .....helpers.ptx_helpers import ( + cp_async_bulk_s2g, + nanosleep, + stg_b64, + stg_f32, + tma_load_1d, +) from .....helpers.smem_workspace import SmemWorkspace from .....helpers.software_sync import SoftwareGridSync from .....helpers.utils import ceil_div, round_up @@ -87,16 +91,12 @@ def from_packed(cls, packed: Union[Int64, Int32]) -> "_SortedElement": return cls(flat_topk_index=Int32(packed), topk_score=None) scratch = cute.make_rmem_tensor((2,), cutlass.Int32) cute.recast_tensor(scratch, cutlass.Int64)[0] = packed - return cls( - flat_topk_index=scratch[0], topk_score=cute.recast_tensor(scratch, cutlass.Float32)[1] - ) + return cls(flat_topk_index=scratch[0], topk_score=cute.recast_tensor(scratch, cutlass.Float32)[1]) @cute.jit def _copy_atom(dtype, num_bits_per_copy: int): - return cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), dtype, num_bits_per_copy=num_bits_per_copy - ) + return cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), dtype, num_bits_per_copy=num_bits_per_copy) @cute.jit @@ -113,9 +113,7 @@ def _smem_exclusive_prefix( num_elements = cute.size(input_tensor) scan_rows = num_elements // 4 num_warps = block_thread_count // 32 - input_vectors = cute.make_tensor( - input_tensor.iterator, cute.make_layout((scan_rows, 4), stride=(4, 1)) - ) + input_vectors = cute.make_tensor(input_tensor.iterator, cute.make_layout((scan_rows, 4), stride=(4, 1))) load_atom = _copy_atom(cutlass.Int32, 128) values = cute.make_rmem_tensor((4,), cutlass.Int32) @@ -129,12 +127,7 @@ def _smem_exclusive_prefix( for element in cutlass.range_constexpr(4): values[element] = Int32(0) - local_prefix = ( - Int32(0), - values[0], - values[0] + values[1], - values[0] + values[1] + values[2], - ) + local_prefix = (Int32(0), values[0], values[0] + values[1], values[0] + values[1] + values[2]) lane_total = local_prefix[3] + values[3] inclusive = lane_total for step_log in cutlass.range_constexpr(5): @@ -231,16 +224,13 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self._validate_router_configuration() self.expert_count_padded = round_up(self.expert_count, 4) self.expert_count_with_trash = self.expert_count_padded + 1 - self.router_elements_per_lane, self.router_data_cta_count = ( - self._router_launch_configuration() - ) + self.router_elements_per_lane, self.router_data_cta_count = self._router_launch_configuration() self.router_tokens_per_cta = self.router_elements_per_lane * self.router_warps_per_cta * 32 self.router_push_cta_count = ceil_div(self.expert_count, self.router_warps_per_cta) self.router_grid_cta_count = max(self.router_data_cta_count + 1, self.router_push_cta_count) if self.router_grid_cta_count > self.promised_launchable_sm_count: raise ValueError( - "Router grid exceeds promised_launchable_sm_count; " - "all metadata-push CTAs must be concurrently resident." + "Router grid exceeds promised_launchable_sm_count; all metadata-push CTAs must be concurrently resident." ) self._router_smem_workspace = self._build_router_smem_workspace() @@ -267,18 +257,14 @@ def _validate_router_configuration(self) -> None: if self.expert_count > 16384: raise NotImplementedError("Local routing supports at most 16384 experts.") if self.topk > self.expert_count: - raise ValueError( - f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}." - ) + raise ValueError(f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}.") def worst_case_padded_tokens(self, block: int) -> int: source_token_capacity = self.max_tokens routes_per_source_token = min(self.topk, self.expert_count) route_capacity = source_token_capacity * routes_per_source_token active_expert_capacity = min(self.expert_count, route_capacity) - route_budget_blocks = ( - active_expert_capacity + (route_capacity - active_expert_capacity) // block - ) + route_budget_blocks = active_expert_capacity + (route_capacity - active_expert_capacity) // block expert_bound_blocks = active_expert_capacity * int(ceil_div(source_token_capacity, block)) return min(route_budget_blocks, expert_bound_blocks) * block @@ -296,13 +282,8 @@ def next_power_of_two(value: int) -> int: maximum_data_cta_count = 128 maximum_supported_tokens = maximum_cta_capacity * maximum_data_cta_count if routed_token_capacity > maximum_supported_tokens: - raise NotImplementedError( - f"The router supports at most {maximum_supported_tokens} routed tokens." - ) - cta_capacity = min( - maximum_cta_capacity, - next_power_of_two(max(routed_token_capacity, minimum_cta_capacity)), - ) + raise NotImplementedError(f"The router supports at most {maximum_supported_tokens} routed tokens.") + cta_capacity = min(maximum_cta_capacity, next_power_of_two(max(routed_token_capacity, minimum_cta_capacity))) elements_per_lane = cta_capacity // (self.router_warps_per_cta * 32) data_cta_count = ceil_div(routed_token_capacity, cta_capacity) return elements_per_lane, data_cta_count @@ -311,18 +292,11 @@ def _build_router_smem_workspace(self) -> SmemWorkspace: workspace = SmemWorkspace() overlay = workspace.create_overlay("local_mega.routing.router_smem.role") data_lifetime = overlay.add_lifetime("data_cta") + data_lifetime.register_tensor(self.router_data_histogram_region, cutlass.Int32, (self.expert_count_with_trash,)) data_lifetime.register_tensor( - self.router_data_histogram_region, cutlass.Int32, (self.expert_count_with_trash,) - ) - data_lifetime.register_tensor( - self.router_data_prefix_region, - cutlass.Int32, - (self.expert_count_with_trash,), - byte_alignment=16, - ) - data_lifetime.register_tensor( - self.router_data_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,) + self.router_data_prefix_region, cutlass.Int32, (self.expert_count_with_trash,), byte_alignment=16 ) + data_lifetime.register_tensor(self.router_data_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,)) data_lifetime.register_tensor( self.router_data_sorted_region, (cutlass.Int64 if self.apply_topk_at_fc1 else cutlass.Int32), @@ -331,30 +305,18 @@ def _build_router_smem_workspace(self) -> SmemWorkspace: ) if self.router_data_cta_count > 1: data_lifetime.register_tensor( - self.router_data_base_region, - cutlass.Int32, - (self.expert_count_padded,), - byte_alignment=16, + self.router_data_base_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 ) helper_lifetime = overlay.add_lifetime("helper_cta") helper_lifetime.register_tensor( - self.router_helper_totals_region, - cutlass.Int32, - (self.expert_count_padded,), - byte_alignment=16, + self.router_helper_totals_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 ) helper_lifetime.register_tensor( - self.router_helper_prefix_region, - cutlass.Int32, - (self.expert_count_padded,), - byte_alignment=16, + self.router_helper_prefix_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 ) helper_lifetime.register_tensor( - self.router_helper_warp_totals_region, - cutlass.Int32, - (self.router_warps_per_cta,), - byte_alignment=16, + self.router_helper_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,), byte_alignment=16 ) workspace.finalize(max_bytes=self.router_smem_limit_bytes) return workspace @@ -370,34 +332,14 @@ def register_device_workspace(self, workspace: DeviceWorkspace) -> None: def _register_router_workspace(self, workspace: DeviceWorkspace) -> None: maximum_routed_tokens = self.max_tokens * self.topk workspace.register( - self.sizes_region, - cutlass.Int32, - (self.expert_count_padded,), - buffer_space="shared", - reset="tail_reset", - ) - workspace.register( - self.sizes_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset" - ) - workspace.register( - self.metadata_ready_region, - cutlass.Int32, - (1,), - buffer_space="shared", - reset="tail_reset", - ) - workspace.register( - self.sorted_metadata_region, - cutlass.Int64, - (maximum_routed_tokens,), - buffer_space="local", + self.sizes_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="shared", reset="tail_reset" ) + workspace.register(self.sizes_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.metadata_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.sorted_metadata_region, cutlass.Int64, (maximum_routed_tokens,), buffer_space="local") if self.apply_topk_at_fc1: workspace.register( - self.sorted_scores_region, - cutlass.Float32, - (maximum_routed_tokens,), - buffer_space="local", + self.sorted_scores_region, cutlass.Float32, (maximum_routed_tokens,), buffer_space="local" ) workspace.register( self.token_src_metadata_region, @@ -408,34 +350,16 @@ def _register_router_workspace(self, workspace: DeviceWorkspace) -> None: ) if self.apply_topk_at_fc1: workspace.register( - self.fc1_topk_scores_region, - cutlass.Float32, - (self.worst_case_token_count,), - buffer_space="shared", + self.fc1_topk_scores_region, cutlass.Float32, (self.worst_case_token_count,), buffer_space="shared" ) + workspace.register(self.pool_expert_base_region, cutlass.Int32, (self.expert_count,), buffer_space="local") workspace.register( - self.pool_expert_base_region, cutlass.Int32, (self.expert_count,), buffer_space="local" + self.source_expert_base_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="local" ) workspace.register( - self.source_expert_base_region, - cutlass.Int32, - (self.expert_count_padded,), - buffer_space="local", - ) - workspace.register( - self.sorted_metadata_ready_region, - cutlass.Int32, - (1,), - buffer_space="local", - reset="tail_reset", - ) - workspace.register( - self.push_table_ready_region, - cutlass.Int32, - (1,), - buffer_space="local", - reset="tail_reset", + self.sorted_metadata_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" ) + workspace.register(self.push_table_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset") if self.router_data_cta_count > 1: workspace.register( self.router_size_counter_region, @@ -445,18 +369,10 @@ def _register_router_workspace(self, workspace: DeviceWorkspace) -> None: reset="tail_reset", ) workspace.register( - self.router_histogram_done_region, - cutlass.Int32, - (1,), - buffer_space="local", - reset="tail_reset", + self.router_histogram_done_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" ) workspace.register( - self.source_base_ready_region, - cutlass.Int32, - (1,), - buffer_space="local", - reset="tail_reset", + self.source_base_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" ) def __extract_mlir_values__(self) -> list: @@ -545,16 +461,10 @@ def _router_single_cta( if self._router_linear_cta_idx < Int32(self.router_data_cta_count): block_thread_count = self.router_warps_per_cta * 32 trash_bucket = self.expert_count_padded - histogram = self._router_smem_workspace.tensor( - self.router_data_histogram_region, smem_base - ) + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) - warp_totals = self._router_smem_workspace.tensor( - self.router_data_warp_totals_region, smem_base - ) - sorted_elements = self._router_smem_workspace.tensor( - self.router_data_sorted_region, smem_base - ) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) zero_round_count = ceil_div(self.expert_count_with_trash, block_thread_count) for zero_round in cutlass.range_constexpr(zero_round_count): @@ -595,12 +505,7 @@ def _router_single_cta( iket.range_push("router.sort") self._sort_router_elements( - expert_registers, - within_expert_indices, - score_registers, - sorted_elements, - prefix, - topk_indices.dtype, + expert_registers, within_expert_indices, score_registers, sorted_elements, prefix, topk_indices.dtype ) cute.arch.sync_threads() iket.range_pop() @@ -609,10 +514,7 @@ def _router_single_cta( cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.sorted_metadata_ready_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" ) iket.range_pop() elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): @@ -626,16 +528,10 @@ def _router_multiple_ctas( block_thread_count = self.router_warps_per_cta * 32 trash_bucket = self.expert_count_padded expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) - histogram = self._router_smem_workspace.tensor( - self.router_data_histogram_region, smem_base - ) + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) - warp_totals = self._router_smem_workspace.tensor( - self.router_data_warp_totals_region, smem_base - ) - sorted_elements = self._router_smem_workspace.tensor( - self.router_data_sorted_region, smem_base - ) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) dump_base = self._router_smem_workspace.tensor(self.router_data_base_region, smem_base) zero_round_count = ceil_div(self.expert_count_with_trash, block_thread_count) @@ -656,17 +552,12 @@ def _router_multiple_ctas( expert = Int32(expert_round * block_thread_count) + self._router_thread_idx if expert < Int32(self.expert_count_padded): dump_base[expert] = Int32( - cute.arch.atomic_add( - size_counter + expert, histogram[expert], sem="relaxed", scope="gpu" - ) + cute.arch.atomic_add(size_counter + expert, histogram[expert], sem="relaxed", scope="gpu") ) cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.router_histogram_done_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.router_histogram_done_region), Int32(1), sem="release", scope="gpu" ) total_valid_routes = _smem_exclusive_prefix( @@ -684,12 +575,7 @@ def _router_multiple_ctas( iket.range_pop() iket.range_push("router.sort") self._sort_router_elements( - expert_registers, - within_expert_indices, - score_registers, - sorted_elements, - prefix, - topk_indices.dtype, + expert_registers, within_expert_indices, score_registers, sorted_elements, prefix, topk_indices.dtype ) cute.arch.sync_threads() iket.range_pop() @@ -697,9 +583,7 @@ def _router_multiple_ctas( iket.range_push("router.wait_source_base") source_base_ready = self._device_workspace.ptr(self.source_base_ready_region) if self._router_thread_idx == Int32(0): - while cute.arch.load(source_base_ready, Int32, sem="acquire", scope="gpu") != Int32( - 1 - ): + while cute.arch.load(source_base_ready, Int32, sem="acquire", scope="gpu") != Int32(1): nanosleep(150) cute.arch.sync_threads() iket.range_pop() @@ -715,10 +599,7 @@ def _router_multiple_ctas( cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.sorted_metadata_ready_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" ) iket.range_pop() elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): @@ -735,9 +616,7 @@ def _router_helper_multiple_ctas(self, smem_base: cute.Pointer) -> None: block_thread_count = self.router_warps_per_cta * 32 totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) - warp_totals = self._router_smem_workspace.tensor( - self.router_helper_warp_totals_region, smem_base - ) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) size_counter = self._device_workspace.tensor(self.router_size_counter_region) source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) @@ -776,10 +655,7 @@ def _router_helper_multiple_ctas(self, smem_base: cute.Pointer) -> None: cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.source_base_ready_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.source_base_ready_region), Int32(1), sem="release", scope="gpu" ) publish_sizes() iket.range_pop() @@ -801,9 +677,7 @@ def _router_push_metadata(self) -> None: nanosleep(150) cute.arch.sync_threads() - global_expert = ( - self._router_linear_cta_idx * Int32(self.router_warps_per_cta) + self._router_warp_idx - ) + global_expert = self._router_linear_cta_idx * Int32(self.router_warps_per_cta) + self._router_warp_idx if global_expert < Int32(self.expert_count): sizes = self._device_workspace.tensor(self.sizes_region) source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) @@ -812,14 +686,10 @@ def _router_push_metadata(self) -> None: source_begin = source_expert_base[global_expert] destination_begin = pool_expert_base[global_expert] source_metadata = self._device_workspace.ptr(self.sorted_metadata_region) - destination_metadata_address = self._device_workspace.ptr( - self.token_src_metadata_region - ).toint() + destination_metadata_address = self._device_workspace.ptr(self.token_src_metadata_region).toint() if cutlass.const_expr(self.apply_topk_at_fc1): source_scores = self._device_workspace.ptr(self.sorted_scores_region) - destination_scores_address = self._device_workspace.ptr( - self.fc1_topk_scores_region - ).toint() + destination_scores_address = self._device_workspace.ptr(self.fc1_topk_scores_region).toint() route_round_count = (route_count + Int32(31)) // Int32(32) for route_round in cutlass.range(route_round_count, unroll=1): route = Int32(route_round) * Int32(32) + self._router_lane_idx @@ -828,16 +698,12 @@ def _router_push_metadata(self) -> None: destination_position = destination_begin + route metadata = cute.arch.load(source_metadata + source_position, cutlass.Int64) stg_b64( - destination_metadata_address - + Int64(destination_position) * Int64(RoutingMetadata.nbytes), + destination_metadata_address + Int64(destination_position) * Int64(RoutingMetadata.nbytes), metadata, ) if cutlass.const_expr(self.apply_topk_at_fc1): score = cute.arch.load(source_scores + source_position, cutlass.Float32) - stg_f32( - destination_scores_address + Int64(destination_position) * Int64(4), - score, - ) + stg_f32(destination_scores_address + Int64(destination_position) * Int64(4), score) cute.arch.sync_threads() if self._router_thread_idx == Int32(0): @@ -846,10 +712,7 @@ def _router_push_metadata(self) -> None: cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.metadata_ready_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.metadata_ready_region), Int32(1), sem="release", scope="gpu" ) @cute.jit @@ -858,9 +721,7 @@ def _publish_sizes(self, smem_expert_counts: cute.Tensor) -> Callable[[], None]: if self._router_warp_idx == Int32(0): with cute.arch.elect_one(): cp_async_bulk_s2g( - self._device_workspace.ptr(self.sizes_region), - smem_expert_counts.iterator, - row_bytes, + self._device_workspace.ptr(self.sizes_region), smem_expert_counts.iterator, row_bytes ) cute.arch.cp_async_bulk_commit_group() @@ -872,10 +733,7 @@ def finalize() -> None: cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.sizes_ready_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.sizes_ready_region), Int32(1), sem="release", scope="gpu" ) return finalize @@ -884,13 +742,9 @@ def finalize() -> None: def _compute_local_tables(self, smem_base: cute.Pointer) -> None: block_thread_count = self.router_warps_per_cta * 32 - padded_totals = self._router_smem_workspace.tensor( - self.router_helper_totals_region, smem_base - ) + padded_totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) - warp_totals = self._router_smem_workspace.tensor( - self.router_helper_warp_totals_region, smem_base - ) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) sizes = self._device_workspace.tensor(self.sizes_region) sizes_ready = self._device_workspace.ptr(self.sizes_ready_region) @@ -908,8 +762,7 @@ def _compute_local_tables(self, smem_base: cute.Pointer) -> None: if expert < Int32(self.expert_count_padded): expert_size = sizes[expert] padded_totals[expert] = ( - (expert_size + Int32(self.token_padding_block - 1)) - // Int32(self.token_padding_block) + (expert_size + Int32(self.token_padding_block - 1)) // Int32(self.token_padding_block) ) * Int32(self.token_padding_block) cute.arch.sync_threads() _smem_exclusive_prefix( @@ -929,10 +782,7 @@ def _compute_local_tables(self, smem_base: cute.Pointer) -> None: cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.push_table_ready_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.push_table_ready_region), Int32(1), sem="release", scope="gpu" ) iket.range_pop() @@ -950,62 +800,46 @@ def _load_router_inputs( elements_per_thread = load_round_count * elements_per_vector topk_flat = cute.make_tensor(topk_indices.iterator, cute.make_layout((maximum_elements,))) - topk_vectors = cute.logical_divide( - cute.zipped_divide(topk_flat, (tile_span,)), (elements_per_vector, None) - ) + topk_vectors = cute.logical_divide(cute.zipped_divide(topk_flat, (tile_span,)), (elements_per_vector, None)) load_atom = _copy_atom(topk_indices.dtype, 128) expert_registers = cute.make_rmem_tensor((elements_per_thread,), cutlass.Int32) if cutlass.const_expr(topk_indices.dtype.width == 64): raw_indices = cute.make_rmem_tensor((elements_per_thread,), topk_indices.dtype) raw_vectors = cute.zipped_divide(raw_indices, (elements_per_vector,)) for load_round in cutlass.range_constexpr(load_round_count): - tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32( - elements_per_vector - ) + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) if tile_begin < actual_elements: cute.copy( load_atom, - mark_alignment( - topk_vectors[(None, self._router_grid_thread_idx), load_round], 16 - ), + mark_alignment(topk_vectors[(None, self._router_grid_thread_idx), load_round], 16), raw_vectors[None, load_round], ) else: expert_vectors = cute.zipped_divide(expert_registers, (elements_per_vector,)) for load_round in cutlass.range_constexpr(load_round_count): - tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32( - elements_per_vector - ) + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) if tile_begin < actual_elements: cute.copy( load_atom, - mark_alignment( - topk_vectors[(None, self._router_grid_thread_idx), load_round], 16 - ), + mark_alignment(topk_vectors[(None, self._router_grid_thread_idx), load_round], 16), expert_vectors[None, load_round], ) score_registers = None if cutlass.const_expr(self.apply_topk_at_fc1): score_registers = cute.make_rmem_tensor((elements_per_thread,), cutlass.Float32) - scores_flat = cute.make_tensor( - topk_scores.iterator, cute.make_layout((maximum_elements,)) - ) + scores_flat = cute.make_tensor(topk_scores.iterator, cute.make_layout((maximum_elements,))) score_vectors = cute.logical_divide( cute.zipped_divide(scores_flat, (tile_span,)), (elements_per_vector, None) ) score_atom = _copy_atom(cutlass.Float32, elements_per_vector * 32) score_register_vectors = cute.zipped_divide(score_registers, (elements_per_vector,)) for load_round in cutlass.range_constexpr(load_round_count): - tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32( - elements_per_vector - ) + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) if tile_begin < actual_elements: cute.copy( score_atom, - mark_alignment( - score_vectors[(None, self._router_grid_thread_idx), load_round], 16 - ), + mark_alignment(score_vectors[(None, self._router_grid_thread_idx), load_round], 16), score_register_vectors[None, load_round], ) @@ -1016,26 +850,21 @@ def _load_router_inputs( if cutlass.const_expr(topk_indices.dtype.width == 64): expert_registers[register_idx] = raw_indices_i32[2 * register_idx] token_idx, _ = self._router_value_coordinate(register_idx, topk_indices.dtype) - is_invalid = ( - expert_registers_u32[register_idx] >= cutlass.Uint32(self.expert_count) - ) | (token_idx >= actual_token_count) + is_invalid = (expert_registers_u32[register_idx] >= cutlass.Uint32(self.expert_count)) | ( + token_idx >= actual_token_count + ) if is_invalid: expert_registers[register_idx] = Int32(self.expert_count_padded) return expert_registers, score_registers @cute.jit - def _build_histogram( - self, expert_registers: cute.Tensor, histogram: cute.Tensor - ) -> cute.Tensor: + def _build_histogram(self, expert_registers: cute.Tensor, histogram: cute.Tensor) -> cute.Tensor: register_count = cute.size(expert_registers) within_expert_indices = cute.make_rmem_tensor((register_count,), cutlass.Int32) for register_idx in cutlass.range_constexpr(register_count): within_expert_indices[register_idx] = Int32( cute.arch.atomic_add( - histogram.iterator + expert_registers[register_idx], - Int32(1), - sem="relaxed", - scope="cta", + histogram.iterator + expert_registers[register_idx], Int32(1), sem="relaxed", scope="cta" ) ) cute.arch.sync_threads() @@ -1055,41 +884,28 @@ def _sort_router_elements( for register_idx in cutlass.range_constexpr(register_count): token_idx, topk_slot = self._router_value_coordinate(register_idx, topk_index_type) flat_topk_index = token_idx * Int32(self.topk) + topk_slot - destination = ( - expert_run_starts[expert_registers[register_idx]] - + within_expert_indices[register_idx] - ) + destination = expert_run_starts[expert_registers[register_idx]] + within_expert_indices[register_idx] if cutlass.const_expr(self.apply_topk_at_fc1): - sorted_elements[destination] = _SortedElement( - flat_topk_index, score_registers[register_idx] - ).pack() + sorted_elements[destination] = _SortedElement(flat_topk_index, score_registers[register_idx]).pack() else: sorted_elements[destination] = _SortedElement(flat_topk_index, None).pack() @cute.jit - def _router_value_coordinate( - self, register_idx: int, topk_index_type: type - ) -> Tuple[Int32, Int32]: + def _router_value_coordinate(self, register_idx: int, topk_index_type: type) -> Tuple[Int32, Int32]: elements_per_vector = 128 // topk_index_type.width - tile_span = ( - elements_per_vector * self.router_data_cta_count * self.router_warps_per_cta * 32 - ) + tile_span = elements_per_vector * self.router_data_cta_count * self.router_warps_per_cta * 32 flat_index = Int32( register_idx // elements_per_vector * tile_span + register_idx % elements_per_vector ) + self._router_grid_thread_idx * Int32(elements_per_vector) return (flat_index // Int32(self.topk), flat_index % Int32(self.topk)) @cute.jit - def _dump_contiguous_router_output( - self, sorted_elements: cute.Tensor, total_valid_routes: Int32 - ) -> None: + def _dump_contiguous_router_output(self, sorted_elements: cute.Tensor, total_valid_routes: Int32) -> None: block_thread_count = self.router_warps_per_cta * 32 metadata_address = self._device_workspace.ptr(self.sorted_metadata_region).toint() if cutlass.const_expr(self.apply_topk_at_fc1): score_address = self._device_workspace.ptr(self.sorted_scores_region).toint() - dump_round_count = (total_valid_routes + Int32(block_thread_count - 1)) // Int32( - block_thread_count - ) + dump_round_count = (total_valid_routes + Int32(block_thread_count - 1)) // Int32(block_thread_count) for dump_round in cutlass.range(dump_round_count, unroll=4): position = Int32(dump_round * block_thread_count) + self._router_thread_idx predicate = Int32(position < total_valid_routes) @@ -1098,11 +914,7 @@ def _dump_contiguous_router_output( token=(element.flat_topk_index // Int32(self.topk)), topk_slot=(element.flat_topk_index % Int32(self.topk)), ) - stg_b64( - metadata_address + Int64(position) * Int64(RoutingMetadata.nbytes), - metadata.pack(), - predicate, - ) + stg_b64(metadata_address + Int64(position) * Int64(RoutingMetadata.nbytes), metadata.pack(), predicate) if cutlass.const_expr(self.apply_topk_at_fc1): stg_f32(score_address + Int64(position) * Int64(4), element.topk_score, predicate) @@ -1128,9 +940,7 @@ def _dump_router_output_by_expert( for route_round in cutlass.range(route_round_count, unroll=1): route = Int32(route_round) * Int32(32) + self._router_lane_idx predicate = Int32(route < run_length) - element = _SortedElement.from_packed( - sorted_elements[predicate * (run_begin + route)] - ) + element = _SortedElement.from_packed(sorted_elements[predicate * (run_begin + route)]) output_position = dump_begin + route metadata = RoutingMetadata( token=(element.flat_topk_index // Int32(self.topk)), @@ -1142,11 +952,7 @@ def _dump_router_output_by_expert( predicate, ) if cutlass.const_expr(self.apply_topk_at_fc1): - stg_f32( - score_address + Int64(output_position) * Int64(4), - element.topk_score, - predicate, - ) + stg_f32(score_address + Int64(output_position) * Int64(4), element.topk_score, predicate) @cute.jit def expert_sizes(self, device_workspace: DeviceWorkspace) -> cute.Tensor: @@ -1171,9 +977,7 @@ def token_src_metadata_pointer(self, device_workspace: DeviceWorkspace) -> cute. return device_workspace.ptr(self.token_src_metadata_region) @cute.jit - def wait_for_sizes_ready( - self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000 - ) -> None: + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: thread_idx, _, _ = cute.arch.thread_idx() lane_idx = thread_idx % Int32(32) if lane_idx == Int32(0): @@ -1183,9 +987,7 @@ def wait_for_sizes_ready( cute.arch.sync_warp() @cute.jit - def wait_for_metadata_ready( - self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000 - ) -> None: + def wait_for_metadata_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: thread_idx, _, _ = cute.arch.thread_idx() lane_idx = thread_idx % Int32(32) if lane_idx == Int32(0): @@ -1291,20 +1093,14 @@ def _validate_configuration(self) -> None: if self.expert_count > 16384: raise NotImplementedError("Local routing supports at most 16384 experts.") if self.topk > self.expert_count: - raise ValueError( - f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}." - ) + raise ValueError(f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}.") if not 1 <= self.token_in_flag_batch <= 32: - raise ValueError( - f"token_in_flag_batch must be in [1, 32], got {self.token_in_flag_batch}." - ) + raise ValueError(f"token_in_flag_batch must be in [1, 32], got {self.token_in_flag_batch}.") if self.tokens_per_fc1_ready_slot % self.token_padding_block != 0: raise ValueError("tokens_per_fc1_ready_slot must be divisible by token_padding_block.") element_block = self.activation_sf_vector_size * 4 if self.hidden_size % element_block != 0: - raise ValueError( - f"{self.quant_kind} requires hidden_size divisible by {element_block}." - ) + raise ValueError(f"{self.quant_kind} requires hidden_size divisible by {element_block}.") if self.sf_padding_block % 128 != 0: raise ValueError("sf_padding_block must be a multiple of 128.") @@ -1336,10 +1132,7 @@ def worst_case_sf_token_count(self) -> int: @property def max_fc1_ready_slot_count(self) -> int: - return ( - self._router.worst_case_padded_tokens(self.tokens_per_fc1_ready_slot) - // self.tokens_per_fc1_ready_slot - ) + return self._router.worst_case_padded_tokens(self.tokens_per_fc1_ready_slot) // self.tokens_per_fc1_ready_slot @property def token_src_metadata_region(self) -> str: @@ -1374,9 +1167,7 @@ def expert_sizes(self, device_workspace: DeviceWorkspace) -> cute.Tensor: return self._router.expert_sizes(device_workspace) @cute.jit - def wait_for_sizes_ready( - self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000 - ) -> None: + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: self._router.wait_for_sizes_ready(device_workspace, sleep_cycles) @cute.jit @@ -1439,9 +1230,7 @@ def _register_main_workspace(self, workspace: DeviceWorkspace) -> None: buffer_space="local", reset="tail_reset", ) - activation_sf_element_count = ( - self.worst_case_sf_token_count * self.activation_sf_hidden_padded - ) + activation_sf_element_count = self.worst_case_sf_token_count * self.activation_sf_hidden_padded workspace.register( self.fc1_activation_sf_region, self.activation_sf_dtype, @@ -1470,10 +1259,7 @@ def register_smem_regions(self, workspace: SmemWorkspace) -> None: token_in_lifetime.register_tensor( self.token_in_sf_smem_region, self.activation_sf_dtype, - ( - self.transfer_warp_count, - (self.activation_sf_vector_size, self.activation_sf_hidden_padded), - ), + (self.transfer_warp_count, (self.activation_sf_vector_size, self.activation_sf_hidden_padded)), stride=(self.activation_sf_hidden_padded, (0, 1)), byte_alignment=16, ) @@ -1487,14 +1273,10 @@ def fc1_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> cute.Te layout = tile_atom_to_shape_SF( (self.worst_case_sf_token_count, self.hidden_size, 1), self.activation_sf_vector_size ) - return cute.make_tensor( - device_workspace.ptr(self.fc1_activation_sf_region), cute.select(layout, mode=[0, 1]) - ) + return cute.make_tensor(device_workspace.ptr(self.fc1_activation_sf_region), cute.select(layout, mode=[0, 1])) @cute.jit - def pre_reduced_activation_tensor( - self, device_workspace: DeviceWorkspace - ) -> Optional[cute.Tensor]: + def pre_reduced_activation_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: """The (tokens, topk, hidden) combine staging plane, or None under in-kernel top-k reduction.""" if cutlass.const_expr(self.reduce_topk_in_kernel): return None @@ -1524,20 +1306,14 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No source_sizes = cute.make_tensor(sizes.iterator, cute.make_layout((self.expert_count,))) copy_elements = 4 if self.expert_count % 4 == 0 else 1 source_size_vectors = cute.zipped_divide( - ( - mark_alignment(source_sizes, 16) - if cutlass.const_expr(copy_elements == 4) - else source_sizes - ), + (mark_alignment(source_sizes, 16) if cutlass.const_expr(copy_elements == 4) else source_sizes), (copy_elements,), ) destination_size_vectors = cute.zipped_divide(owned_sizes, (copy_elements,)) size_vector_count = cute.size(destination_size_vectors, mode=[1]) size_copy_atom = cute.make_copy_atom( cute.nvgpu.cpasync.CopyG2SOp( - cache_mode=cute.nvgpu.LoadCacheMode.GLOBAL - if copy_elements == 4 - else cute.nvgpu.LoadCacheMode.ALWAYS + cache_mode=cute.nvgpu.LoadCacheMode.GLOBAL if copy_elements == 4 else cute.nvgpu.LoadCacheMode.ALWAYS ), cutlass.Int32, num_bits_per_copy=copy_elements * 32, @@ -1548,9 +1324,7 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No vector_idx = Int32(size_copy_round * self.transfer_thread_count) + transfer_thread_idx if vector_idx < Int32(size_vector_count): cute.copy( - size_copy_atom, - source_size_vectors[None, vector_idx], - destination_size_vectors[None, vector_idx], + size_copy_atom, source_size_vectors[None, vector_idx], destination_size_vectors[None, vector_idx] ) cute.arch.cp_async_commit_group() iket.range_pop() @@ -1582,9 +1356,7 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No sf_copy_elements = 4 source_sf_values = cute.slice_(warp_sf_stage, (0, None)) source_sf_vectors = cute.zipped_divide(source_sf_values, (sf_copy_elements,)) - sf_copy_atom = _copy_atom( - self.activation_sf_dtype, sf_copy_elements * int(self.activation_sf_dtype.width) - ) + sf_copy_atom = _copy_atom(self.activation_sf_dtype, sf_copy_elements * int(self.activation_sf_dtype.width)) next_dense_token = global_warp_idx expert_valid_begin = Int32(0) @@ -1605,34 +1377,24 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No expert_valid_end = expert_valid_begin + expert_token_count pull_count = Int32(0) if next_dense_token < expert_valid_end: - pull_count = ( - expert_valid_end - next_dense_token + global_warp_count - Int32(1) - ) // global_warp_count + pull_count = (expert_valid_end - next_dense_token + global_warp_count - Int32(1)) // global_warp_count for pull_round in cutlass.range(pull_count, unroll=1): dense_token_idx = next_dense_token + Int32(pull_round) * global_warp_count token_in_expert = dense_token_idx - expert_valid_begin pool_token_idx = pool_expert_bases[local_expert] + token_in_expert sf_token_idx = expert_sf_begin + token_in_expert - ready_slot_idx = expert_ready_slot_begin + token_in_expert // Int32( - self.tokens_per_fc1_ready_slot - ) + ready_slot_idx = expert_ready_slot_begin + token_in_expert // Int32(self.tokens_per_fc1_ready_slot) metadata = RoutingMetadata.load( - token_metadata_pointer.toint() - + Int64(pool_token_idx) * Int64(RoutingMetadata.nbytes) + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(RoutingMetadata.nbytes) ) source_sf_row = self._local_routing_args.activation_sf[Int64(metadata.token), None] with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - warp_mbarrier, Int32(activation_sf_bytes) - ) + cute.arch.mbarrier_arrive_and_expect_tx(warp_mbarrier, Int32(activation_sf_bytes)) tma_load_1d( - warp_sf_stage.iterator, - source_sf_row.iterator, - warp_mbarrier, - Int32(activation_sf_bytes), + warp_sf_stage.iterator, source_sf_row.iterator, warp_mbarrier, Int32(activation_sf_bytes) ) cute.arch.sync_warp() cute.arch.mbarrier_wait(warp_mbarrier, pull_phase) @@ -1640,9 +1402,7 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No destination_sf_row = fc1_activation_sf[Int64(sf_token_idx), ((None, None), None)] destination_sf_values = cute.slice_(destination_sf_row, (0, None, None)) destination_sf_values = cute.group_modes(destination_sf_values, 0, 2) - destination_sf_vectors = cute.zipped_divide( - destination_sf_values, (sf_copy_elements,) - ) + destination_sf_vectors = cute.zipped_divide(destination_sf_values, (sf_copy_elements,)) sf_vector_count = cute.size(destination_sf_vectors, mode=[1]) for sf_round in cutlass.range_constexpr(ceil_div(sf_vector_count, 32)): sf_vector_idx = Int32(sf_round * 32) + lane_idx @@ -1655,17 +1415,14 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No cute.arch.sync_warp() ready_address = (fc1_ready_counter + ready_slot_idx).toint() - flag_tracker = flag_tracker.accumulate( - Int32(0), self.token_in_flag_batch, ready_address - ) + flag_tracker = flag_tracker.accumulate(Int32(0), self.token_in_flag_batch, ready_address) cute.arch.sync_warp() pull_phase = pull_phase ^ Int32(1) next_dense_token = next_dense_token + pull_count * global_warp_count expert_valid_begin = expert_valid_end expert_sf_begin = expert_sf_begin + ( - (expert_token_count + Int32(self.sf_padding_block - 1)) - // Int32(self.sf_padding_block) + (expert_token_count + Int32(self.sf_padding_block - 1)) // Int32(self.sf_padding_block) ) * Int32(self.sf_padding_block) expert_ready_slot_begin = expert_ready_slot_begin + ( (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) @@ -1692,9 +1449,7 @@ def reset_tail(self) -> None: ) iket.range_pop() total_reset_threads = self.promised_launchable_sm_count * self.transfer_thread_count - global_reset_thread = ( - self._linear_cta_idx * Int32(self.transfer_thread_count) + transfer_thread_idx - ) + global_reset_thread = self._linear_cta_idx * Int32(self.transfer_thread_count) + transfer_thread_idx iket.range_push("tail.reset_workspace") self._device_workspace.reset_tail_space("shared", global_reset_thread, total_reset_threads) self._device_workspace.reset_tail_space("local", global_reset_thread, total_reset_threads) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/tma_gather.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/tma_gather.py index 4e26bd3d722c..b96869775c12 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/tma_gather.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/tma_gather.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Rubin TMA Gather4 load primitives used by Local MegaMoE.""" import cutlass @@ -10,6 +7,7 @@ from cutlass.cute.typing import AddressSpace from cutlass.cutlass_dsl import dsl_user_op + _CTA_GROUP_2_MBAR_MASK = 0xFEFFFFFF diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py index db92666ff682..f0868f743ba4 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py @@ -10,29 +10,26 @@ """ import dataclasses +import math from typing import Any, Callable, ClassVar, List, Literal, Optional, Tuple, Type, Union import cutlass import cutlass.cute as cute -import cutlass.pipeline as pipeline +from cutlass.cutlass_dsl import Int64, T +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.typing import AddressSpace import cutlass.utils as utils +import cutlass.pipeline as pipeline import cutlass.utils.blackwell_helpers as sm100_utils + from cutlass._mlir import ir from cutlass._mlir.dialects import llvm, vector -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass.cute.typing import AddressSpace -from cutlass.cutlass_dsl import Int64, T -from .....api import ( - ImplDesc, - KernelComponent, - OptionalRequirement, - ProblemDesc, - StaticOrRuntimeIntegerType, -) from .....communication.nvlink_domain.symmetric_buffer import SymmetricBufferDevice from .....communication.token_protocol import TokenSrcMetadata -from .....helpers.constants import Fp8E4M3RcpLimit, Fp8E5M2RcpLimit, Fp32Max, Nvfp4E2M1RcpLimit +from .....quant_def import CombineFormat, QuantKind +from .....api import ImplDesc, KernelComponent, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType +from .....helpers.constants import Nvfp4E2M1RcpLimit, Fp8E4M3RcpLimit, Fp8E5M2RcpLimit, Fp32Max from .....helpers.cute_py_helpers import tcgen05_block_scaled_acc_dtype from .....helpers.dsl_helpers import mark_alignment from .....helpers.flag_batch import make_flag_batch_tracker @@ -44,10 +41,8 @@ red_add_relaxed_sys_v2_bf16x2, ) from .....helpers.smem_workspace import SmemRegion, SmemWorkspace -from .....quant_def import CombineFormat, QuantKind from ....function_mapping import CoordinateSpace, FunctionMapping -from ....schedulers.base import SchedulerConsumer -from ....schedulers.fc12_mapping import BlockPhase, SwapAbFc12WorkTileInfo +from ....schedulers import BlockPhase, SchedulerConsumer, SwapAbFc12WorkTileInfo from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension @@ -82,43 +77,27 @@ class QuantImpl: warp_idx: Optional[Any] = None pair_exchange_barrier: Optional[Any] = None - _directions: ClassVar[Tuple[str, ...]] = ( - "regs_in_thread", - "threads_with_the_same_reg", - "regs_in_pair_threads", - ) + _directions: ClassVar[Tuple[str, ...]] = ("regs_in_thread", "threads_with_the_same_reg", "regs_in_pair_threads") # -- config / validation -------------------------------------------------- def __post_init__(self): if isinstance(self.quant_kind, CombineFormat): if not self.quant_kind.is_quantized: - raise ValueError( - f"QuantImpl combine path needs a quantized CombineFormat, got {self.quant_kind}." - ) + raise ValueError(f"QuantImpl combine path needs a quantized CombineFormat, got {self.quant_kind}.") elif not isinstance(self.quant_kind, QuantKind): - raise ValueError( - f"quant_kind must be a QuantKind or a CombineFormat, got {self.quant_kind!r}." - ) + raise ValueError(f"quant_kind must be a QuantKind or a CombineFormat, got {self.quant_kind!r}.") if self.sf_vec_direction not in self._directions: - raise ValueError( - f"sf_vec_direction must be one of {self._directions}, got {self.sf_vec_direction!r}." - ) + raise ValueError(f"sf_vec_direction must be one of {self._directions}, got {self.sf_vec_direction!r}.") if self.sf_vec_direction == "threads_with_the_same_reg" and self.lane_idx is None: raise ValueError("across-lane quant needs lane_idx for the CREDUX half-warp predicate.") if self.sf_vec_direction == "regs_in_pair_threads": if isinstance(self.quant_kind, CombineFormat): - raise ValueError( - "The paired direction is orthodox-only; combine blocks never straddle warps." - ) + raise ValueError("The paired direction is orthodox-only; combine blocks never straddle warps.") if self.sf_vec_size % 2 != 0: - raise ValueError( - f"The paired direction needs an even sf_vec, got {self.sf_vec_size}." - ) + raise ValueError(f"The paired direction needs an even sf_vec, got {self.sf_vec_size}.") if self.lane_idx is None or self.warp_idx is None or self.pair_exchange_barrier is None: - raise ValueError( - "The paired direction needs lane_idx, warp_idx and pair_exchange_barrier." - ) + raise ValueError("The paired direction needs lane_idx, warp_idx and pair_exchange_barrier.") elif not isinstance(self.quant_kind, CombineFormat) and self.sf_vec_size != 16: # A wider orthodox block cannot be reduced inside one thread: see the paired direction. raise ValueError(f"Orthodox sf_vec {self.sf_vec_size} needs the paired direction.") @@ -159,20 +138,14 @@ def _data_rcp_limit(self) -> float: @property def _block_share_per_thread(self) -> int: """How many of a block's elements this thread holds.""" - return ( - self.sf_vec_size // 2 - if self.sf_vec_direction == "regs_in_pair_threads" - else self.sf_vec_size - ) + return self.sf_vec_size // 2 if self.sf_vec_direction == "regs_in_pair_threads" else self.sf_vec_size # -- dispatch ------------------------------------------------------------- @cute.jit def __call__(self, prequant_reg: cute.Tensor, *, norm_const=None, smem_intermediate=None): if cutlass.const_expr(cute.size(prequant_reg) % self._block_share_per_thread != 0): - raise ValueError( - "prequant_reg size must be divisible by this thread's share of a block." - ) + raise ValueError("prequant_reg size must be divisible by this thread's share of a block.") # Combine quantizes fc2's bf16 reorder regs; orthodox the fp32 swiglu. expected_in = cutlass.BFloat16 if self._is_combine else cutlass.Float32 if cutlass.const_expr(prequant_reg.element_type is not expected_in): @@ -181,9 +154,7 @@ def __call__(self, prequant_reg: cute.Tensor, *, norm_const=None, smem_intermedi ) needs_smem = cutlass.const_expr(self.sf_vec_direction == "regs_in_pair_threads") if cutlass.const_expr(needs_smem != (smem_intermediate is not None)): - raise ValueError( - "smem_intermediate must be supplied for the paired direction and only for it." - ) + raise ValueError("smem_intermediate must be supplied for the paired direction and only for it.") if cutlass.const_expr(not self._is_combine): if cutlass.const_expr(self.sf_vec_direction == "regs_in_pair_threads"): return self.mx_quant_regs_in_pair_threads_impl(prequant_reg, smem_intermediate) @@ -255,14 +226,10 @@ def mx_quant_regs_in_pair_threads_impl( if cutlass.const_expr(smem_intermediate.stride[0] != 1): # Thread-major would still be correct but would silently split the exchange into # per-block scalar accesses. - raise ValueError( - "The paired amax exchange must be block-major to stay one vector access." - ) + raise ValueError("The paired amax exchange must be block-major to stay one vector access.") exchange_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), - cutlass.Float32, - num_bits_per_copy=n_blocks * cutlass.Float32.width, + cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=n_blocks * cutlass.Float32.width ) # Which warps pair up is epilogue knowledge: warp w owns one contiguous run of # intermediate outputs, so a block that is twice that run wide joins adjacent warps. @@ -273,20 +240,12 @@ def mx_quant_regs_in_pair_threads_impl( half_amax = cute.make_rmem_tensor((n_blocks,), cutlass.Float32) for vec_block_idx in cutlass.range_constexpr(n_blocks): half_amax[vec_block_idx] = self._amax_thread_fp32(in_blocks[None, vec_block_idx]) - cute.copy( - exchange_atom, - cute.coalesce(half_amax), - cute.coalesce(smem_intermediate[None, own_thread]), - ) + cute.copy(exchange_atom, cute.coalesce(half_amax), cute.coalesce(smem_intermediate[None, own_thread])) self.pair_exchange_barrier.arrive_and_wait() partner_amax = cute.make_rmem_tensor((n_blocks,), cutlass.Float32) - cute.copy( - exchange_atom, - cute.coalesce(smem_intermediate[None, partner_thread]), - cute.coalesce(partner_amax), - ) + cute.copy(exchange_atom, cute.coalesce(smem_intermediate[None, partner_thread]), cute.coalesce(partner_amax)) data_blocks = [] sf_values = [] @@ -294,17 +253,13 @@ def mx_quant_regs_in_pair_threads_impl( block_amax = cute.arch.fmax(half_amax[vec_block_idx], partner_amax[vec_block_idx]) scale_e8m0, scale_f32 = self._e8m0(block_amax) sf_values.append(scale_e8m0) - data_blocks.append( - self._scale_to_data_ssa(in_blocks[None, vec_block_idx], self._enc_mxfp8(scale_f32)) - ) + data_blocks.append(self._scale_to_data_ssa(in_blocks[None, vec_block_idx], self._enc_mxfp8(scale_f32))) self._store_packed_blocks(data, data_blocks) sf.store(self._values_to_ssa(sf_values, self.scale_dtype)) return data, sf @cute.jit - def nvfp4_combine_quant_regs_in_thread_impl( - self, prequant_reg: cute.Tensor - ) -> Tuple[cute.Tensor, cute.Tensor]: + def nvfp4_combine_quant_regs_in_thread_impl(self, prequant_reg: cute.Tensor) -> Tuple[cute.Tensor, cute.Tensor]: # bf16 in -> e2m1 data + per-16 bf16 amax. amax found on bf16 (packed). vec = self.sf_vec_size n_blocks = cute.size(prequant_reg) // vec @@ -361,9 +316,7 @@ def nvfp4_combine_quant_threads_with_the_same_reg_impl( return data, selected_sf @cute.jit - def mxfp8_combine_quant_regs_in_thread_impl( - self, prequant_reg: cute.Tensor - ) -> Tuple[cute.Tensor, cute.Tensor]: + def mxfp8_combine_quant_regs_in_thread_impl(self, prequant_reg: cute.Tensor) -> Tuple[cute.Tensor, cute.Tensor]: # bf16 in -> e4m3/e5m2 data + per-32 e8m0. amax found on bf16 (packed). vec = self.sf_vec_size n_blocks = cute.size(prequant_reg) // vec @@ -439,8 +392,7 @@ def _concat_blocks_ssa(self, blocks, dtype: Type[cutlass.Numeric]) -> cute.Tenso @cute.jit def _values_to_ssa(self, values, dtype: Type[cutlass.Numeric]) -> cute.TensorSSA: vec = vector.from_elements( - T.vector(len(values), dtype.mlir_type), - [values[i].ir_value() for i in range(len(values))], + T.vector(len(values), dtype.mlir_type), [values[i].ir_value() for i in range(len(values))] ) return cute.TensorSSA(vec, (len(values),), dtype) @@ -484,9 +436,7 @@ def max_abs(lhs: cutlass.Float32, rhs: cutlass.Float32) -> cutlass.Float32: acc = block[0] for elem_idx in cutlass.range_constexpr(1, cute.size(block)): acc = max_abs(acc, block[elem_idx]) - mag_bits = cutlass.Int32( - llvm.bitcast(T.i32(), cutlass.Float32(acc).ir_value()) - ) & cutlass.Int32(0x7FFFFFFF) + mag_bits = cutlass.Int32(llvm.bitcast(T.i32(), cutlass.Float32(acc).ir_value())) & cutlass.Int32(0x7FFFFFFF) return cutlass.Float32(llvm.bitcast(T.f32(), mag_bits.ir_value())) @cute.jit @@ -592,15 +542,12 @@ class _TmemTranspose16x32Core: @staticmethod def _tmem_layout(num_lanes: int, num_cols: int) -> cute.Layout: return cute.make_layout( - (((num_lanes, num_cols), 1),), - stride=(((_TmemTranspose16x32Core._TmemRowStride, 1), 0),), + (((num_lanes, num_cols), 1),), stride=(((_TmemTranspose16x32Core._TmemRowStride, 1), 0),) ) @staticmethod def _rmem_copy_view(rmem: cute.Tensor, num_regs: int, offset: int = 0) -> cute.Tensor: - return cute.make_tensor( - rmem.iterator + offset, cute.make_layout((((num_regs,), 1),), stride=(((1,), 0),)) - ) + return cute.make_tensor(rmem.iterator + offset, cute.make_layout((((num_regs,), 1),), stride=(((1,), 0),))) def __init__(self, tmem_ptr, region: int, reg_tensor: Optional[cute.Tensor] = None) -> None: # The whole transpose is built from 32-bit element atoms; _io_dtype @@ -631,21 +578,11 @@ def __init__(self, tmem_ptr, region: int, reg_tensor: Optional[cute.Tensor] = No self._tmem_dst_top = cute.make_tensor(dst_ptr, self._tmem_layout(16, 16)) self._tmem_dst_bot = cute.make_tensor(dst_ptr + half_lane_off, self._tmem_layout(16, 16)) - self._atom_ld16x64 = cute.make_copy_atom( - tcgen05.Ld16x64bOp(tcgen05.Repetition.x16), self._io_dtype - ) - self._atom_st16x128 = cute.make_copy_atom( - tcgen05.St16x128bOp(tcgen05.Repetition.x8), self._io_dtype - ) - self._atom_st32x32 = cute.make_copy_atom( - tcgen05.St32x32bOp(tcgen05.Repetition.x16), self._io_dtype - ) - self._atom_ld16x256 = cute.make_copy_atom( - tcgen05.Ld16x256bOp(tcgen05.Repetition.x2), self._io_dtype - ) - self._atom_ld16x128 = cute.make_copy_atom( - tcgen05.Ld16x128bOp(tcgen05.Repetition.x4), self._io_dtype - ) + self._atom_ld16x64 = cute.make_copy_atom(tcgen05.Ld16x64bOp(tcgen05.Repetition.x16), self._io_dtype) + self._atom_st16x128 = cute.make_copy_atom(tcgen05.St16x128bOp(tcgen05.Repetition.x8), self._io_dtype) + self._atom_st32x32 = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x16), self._io_dtype) + self._atom_ld16x256 = cute.make_copy_atom(tcgen05.Ld16x256bOp(tcgen05.Repetition.x2), self._io_dtype) + self._atom_ld16x128 = cute.make_copy_atom(tcgen05.Ld16x128bOp(tcgen05.Repetition.x4), self._io_dtype) self._src_regs = cute.make_rmem_tensor((16,), self._io_dtype) # ``output`` is a bare (16,) RMEM fragment; its (lane_idx, elem_idx) @@ -700,18 +637,10 @@ def r2_store(self) -> None: # -- R3 ------------------------------------------------------------------ def r3_load_top(self) -> None: - cute.copy( - self._atom_ld16x256, - self._tmem_dst_top, - self._rmem_copy_view(self._src_regs, 8, offset=0), - ) + cute.copy(self._atom_ld16x256, self._tmem_dst_top, self._rmem_copy_view(self._src_regs, 8, offset=0)) def r3_load_bot(self) -> None: - cute.copy( - self._atom_ld16x256, - self._tmem_dst_bot, - self._rmem_copy_view(self._src_regs, 8, offset=8), - ) + cute.copy(self._atom_ld16x256, self._tmem_dst_bot, self._rmem_copy_view(self._src_regs, 8, offset=8)) def r3_perm(self) -> None: for r in range(16): @@ -723,18 +652,10 @@ def r3_store(self) -> None: # -- R4 ------------------------------------------------------------------ def r4_load_top(self) -> None: - cute.copy( - self._atom_ld16x128, - self._tmem_dst_top, - self._rmem_copy_view(self._src_regs, 8, offset=0), - ) + cute.copy(self._atom_ld16x128, self._tmem_dst_top, self._rmem_copy_view(self._src_regs, 8, offset=0)) def r4_load_bot(self) -> None: - cute.copy( - self._atom_ld16x128, - self._tmem_dst_bot, - self._rmem_copy_view(self._src_regs, 8, offset=8), - ) + cute.copy(self._atom_ld16x128, self._tmem_dst_bot, self._rmem_copy_view(self._src_regs, 8, offset=8)) def r4_perm(self) -> None: for r in range(16): @@ -794,10 +715,7 @@ class TmemTranspose32x32Inplace: """ def __init__( - self, - tmem_ptr, - reg_tensor_top: Optional[cute.Tensor] = None, - reg_tensor_bot: Optional[cute.Tensor] = None, + self, tmem_ptr, reg_tensor_top: Optional[cute.Tensor] = None, reg_tensor_bot: Optional[cute.Tensor] = None ) -> None: if (reg_tensor_top is None) != (reg_tensor_bot is None): raise ValueError( @@ -890,26 +808,19 @@ class TmemTranspose32x64B16Movm: @staticmethod def _tmem_layout(num_lanes: int, num_cols: int) -> cute.Layout: return cute.make_layout( - (((num_lanes, num_cols), 1),), - stride=(((TmemTranspose32x64B16Movm._tmem_row_stride, 1), 0),), + (((num_lanes, num_cols), 1),), stride=(((TmemTranspose32x64B16Movm._tmem_row_stride, 1), 0),) ) @staticmethod def _rmem_copy_view(rmem: cute.Tensor, num_regs: int, offset: int = 0) -> cute.Tensor: - return cute.make_tensor( - rmem.iterator + offset, cute.make_layout((((num_regs,), 1),), stride=(((1,), 0),)) - ) + return cute.make_tensor(rmem.iterator + offset, cute.make_layout((((num_regs,), 1),), stride=(((1,), 0),))) @cute.jit def __init__(self, tmem_ptr, reg_tensor: cute.Tensor) -> None: if cutlass.const_expr(reg_tensor.element_type is not cutlass.BFloat16): - raise TypeError( - f"{type(self).__name__} expects BF16 input after f2fp, got {reg_tensor.element_type}." - ) + raise TypeError(f"{type(self).__name__} expects BF16 input after f2fp, got {reg_tensor.element_type}.") if cutlass.const_expr(cute.size(reg_tensor) != 64): - raise ValueError( - f"{type(self).__name__} expects 64 BF16 elements, got {cute.size(reg_tensor)}." - ) + raise ValueError(f"{type(self).__name__} expects 64 BF16 elements, got {cute.size(reg_tensor)}.") movm_words = movmatrix_b16(cute.recast_tensor(reg_tensor, cutlass.Int32)) self._store_words = cute.make_rmem_tensor(movm_words.layout, movm_words.element_type) @@ -920,20 +831,14 @@ def __init__(self, tmem_ptr, reg_tensor: cute.Tensor) -> None: self._tmem_top = cute.make_tensor(tmem_ptr, self._tmem_layout(16, 32)) self._tmem_bottom = cute.make_tensor(tmem_ptr + half_lane_offset, self._tmem_layout(16, 32)) self._tmem_full = cute.make_tensor(tmem_ptr, self._tmem_layout(32, 32)) - self._store_atom = cute.make_copy_atom( - tcgen05.St16x128bOp(tcgen05.Repetition.x8), cutlass.Float32 - ) - self._load_atom = cute.make_copy_atom( - tcgen05.Ld32x32bOp(tcgen05.Repetition.x32), cutlass.Float32 - ) + self._store_atom = cute.make_copy_atom(tcgen05.St16x128bOp(tcgen05.Repetition.x8), cutlass.Float32) + self._load_atom = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x32), cutlass.Float32) @cute.jit def __call__(self) -> cute.Tensor: movm_words_f32 = cute.recast_tensor(self._store_words, cutlass.Float32) cute.copy(self._store_atom, self._rmem_copy_view(movm_words_f32, 16), self._tmem_top) - cute.copy( - self._store_atom, self._rmem_copy_view(movm_words_f32, 16, offset=16), self._tmem_bottom - ) + cute.copy(self._store_atom, self._rmem_copy_view(movm_words_f32, 16, offset=16), self._tmem_bottom) output_words = cute.make_rmem_tensor((32,), cutlass.Float32) cute.copy(self._load_atom, self._tmem_full, self._rmem_copy_view(output_words, 32)) @@ -948,8 +853,7 @@ class GatedActEpilogueArgs: fc2_alpha: Optional[cute.Tensor] fc1_norm_const: Optional[cute.Tensor] # ----------------------------------- - # MoE domain (token, topk), deepgemm graph only? for transformer graph, we want reduce kernel to perform the - # score mul. + # MoE domain (token, topk), deepgemm graph only? for transformer graph, we want reduce kernel to perform the score mul. topk_scores: Optional[cute.Tensor] @@ -1013,6 +917,7 @@ def impl_desc_require(cls) -> dict[str, object]: "fc2_tma_stages": OptionalRequirement(int), "reduce_topk_in_kernel": OptionalRequirement(bool), "token_back_push_data": OptionalRequirement(bool), + "token_back_ready_granularity": OptionalRequirement(str), } @staticmethod @@ -1027,19 +932,13 @@ def _resolve_situ_betas(problem_desc: ProblemDesc) -> Tuple[Optional[float], Opt beta = problem_desc.get("situ_beta") linear_beta = problem_desc.get("situ_linear_beta") if (beta is None) != (linear_beta is None): - raise ValueError( - f"situ_beta and situ_linear_beta must be set together, got {beta} and {linear_beta}." - ) + raise ValueError(f"situ_beta and situ_linear_beta must be set together, got {beta} and {linear_beta}.") if beta is None: return None, None if beta <= 0 or linear_beta <= 0: - raise ValueError( - f"SiTU beta parameters must be positive, got {beta} and {linear_beta}." - ) + raise ValueError(f"SiTU beta parameters must be positive, got {beta} and {linear_beta}.") if problem_desc["gate_up_clamp"] is not None: - raise ValueError( - "SiTU does not support gate_up_clamp; the two activation variants are exclusive." - ) + raise ValueError("SiTU does not support gate_up_clamp; the two activation variants are exclusive.") if problem_desc.get("swiglu_alpha") is not None: raise ValueError("SiTU does not support SwiGLUBias constants.") return float(beta), float(linear_beta) @@ -1068,21 +967,18 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.communication_enabled = impl_desc["communication_enabled"] self.fc1_epi_flag_batch = impl_desc["fc1_epi_flag_batch"] self.fc2_epi_flag_batch = impl_desc["fc2_epi_flag_batch"] + self.token_back_ready_granularity = impl_desc.get("token_back_ready_granularity", "expert") + if self.token_back_ready_granularity not in ("expert", "token_tile"): + raise ValueError(f"Unsupported token_back_ready_granularity {self.token_back_ready_granularity!r}.") if self.communication_enabled: for field_name in ("reduce_topk_in_kernel", "token_back_push_data"): if field_name not in impl_desc: - raise KeyError( - f"Communication-enabled Epilogue requires ImplDesc field {field_name!r}." - ) + raise KeyError(f"Communication-enabled Epilogue requires ImplDesc field {field_name!r}.") self.reduce_topk_in_kernel = impl_desc.get("reduce_topk_in_kernel", False) self.token_back_push_data = impl_desc.get("token_back_push_data", False) - if not self.communication_enabled and ( - self.reduce_topk_in_kernel or self.token_back_push_data - ): - raise ValueError( - "A communication-disabled Epilogue cannot enable communication policies." - ) + if not self.communication_enabled and (self.reduce_topk_in_kernel or self.token_back_push_data): + raise ValueError("A communication-disabled Epilogue cannot enable communication policies.") # FC1 emits what FC2 consumes as its activation, in the kind's own scale format. self.fc1_output_dtype = self.quant_kind.activation_dtype @@ -1104,12 +1000,8 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: not 1 <= self.fc1_epi_flag_batch <= self._EpilogueWarpCnt or not 1 <= self.fc2_epi_flag_batch <= self._EpilogueWarpCnt ): - raise ValueError( - f"Asynchronous epilogue flag batch sizes must be in [1, {self._EpilogueWarpCnt}]." - ) - self.cluster_tile_intermediate_downproj = ( - self._EpilogueFc1IntermediateDownTileSize * self.cluster_shape_mn[0] - ) + raise ValueError(f"Asynchronous epilogue flag batch sizes must be in [1, {self._EpilogueWarpCnt}].") + self.cluster_tile_intermediate_downproj = self._EpilogueFc1IntermediateDownTileSize * self.cluster_shape_mn[0] atom_thr_size = 2 if self.use_2cta_instrs else 1 self.cta_tile_m = self._EpilogueFc2HiddenTileSize @@ -1127,21 +1019,12 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: if tmem_plan.accumulator_stage_stride_columns != tmem_plan.accumulator_stage_columns: raise ValueError("Rubin MegaMoE does not support overlapping accumulator stages.") if tmem_plan.accumulator_pipeline_stages != tmem_plan.accumulator_stage_count: - raise ValueError( - "Rubin MegaMoE requires one pipeline stage per disjoint accumulator stage." - ) + raise ValueError("Rubin MegaMoE requires one pipeline stage per disjoint accumulator stage.") self.num_accumulator_tmem_cols = tmem_plan.accumulator_columns - self.accumulator_shape = ( - self.cta_tile_m, - self.cta_tile_n, - tmem_plan.accumulator_stage_count, - ) + self.accumulator_shape = (self.cta_tile_m, self.cta_tile_n, tmem_plan.accumulator_stage_count) self.accumulator_stride = (1 << 16, 1, tmem_plan.accumulator_stage_stride_columns) - if ( - isinstance(self.hidden_size, int) - and self.hidden_size % (self.cta_tile_m * self.cluster_shape_mn[0]) == 0 - ): + if isinstance(self.hidden_size, int) and self.hidden_size % (self.cta_tile_m * self.cluster_shape_mn[0]) == 0: self.fc2_hidden_needs_predicate: bool = False else: self.fc2_hidden_needs_predicate: bool = True @@ -1155,10 +1038,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: # One staging stage per token subtile, each an (epi_tile_n, epi_tile_m // 2) quantized tile. self.fc1_staging_stage_bytes = ( - self._EpilogueTokenTileSize - * self._EpilogueFc1IntermediateDownTileSize - * self.fc1_output_dtype.width - // 8 + self._EpilogueTokenTileSize * self._EpilogueFc1IntermediateDownTileSize * self.fc1_output_dtype.width // 8 ) self.fc1_staging_bytes = self.subtile_cnt * self.fc1_staging_stage_bytes # The amax exchange plane is (token chunk, epilogue thread), block-major so each thread's @@ -1167,21 +1047,14 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: # bytes -- see the lifetime argument in fc1_quant. self.fc1_amax_token_chunks = self._EpilogueTokenTileSize // 32 self.fc1_amax_slot_count = ( - self.fc1_amax_token_chunks * self._EpilogueWarpCnt * 32 - if self.needs_pair_amax_exchange - else 0 + self.fc1_amax_token_chunks * self._EpilogueWarpCnt * 32 if self.needs_pair_amax_exchange else 0 ) if self.fc1_amax_slot_count * 4 > self.fc1_staging_stage_bytes: raise ValueError("The paired amax exchange does not fit in one fc1 staging stage.") requested_fc2_tma_stages = impl_desc.get("fc2_tma_stages") - if ( - requested_fc2_tma_stages is not None - and not 1 <= requested_fc2_tma_stages <= self.subtile_cnt - ): - raise ValueError( - f"fc2_tma_stages must be in [1, {self.subtile_cnt}], got {requested_fc2_tma_stages}." - ) + if requested_fc2_tma_stages is not None and not 1 <= requested_fc2_tma_stages <= self.subtile_cnt: + raise ValueError(f"fc2_tma_stages must be in [1, {self.subtile_cnt}], got {requested_fc2_tma_stages}.") if self.fc2_use_bulk: single_stage_region = self._make_fc2_single_stage_region() if self.fc2_use_ublk and single_stage_region.nbytes % 16 != 0: @@ -1191,20 +1064,11 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: else: bf16_baseline_stages = min(2, self.subtile_cnt) bf16_stage_bytes = ( - self._EpilogueTokenTileSize - * self._EpilogueFc2HiddenTileSize - * cutlass.BFloat16.width - // 8 - ) - available_staging_bytes = max( - self.fc1_staging_bytes, bf16_baseline_stages * bf16_stage_bytes - ) - self.fc2_tma_stages = min( - self.subtile_cnt, available_staging_bytes // single_stage_region.nbytes + self._EpilogueTokenTileSize * self._EpilogueFc2HiddenTileSize * cutlass.BFloat16.width // 8 ) - self.fc2_staging_spec: Optional[SmemRegion] = self._make_fc2_staging_region( - self.fc2_tma_stages - ) + available_staging_bytes = max(self.fc1_staging_bytes, bf16_baseline_stages * bf16_stage_bytes) + self.fc2_tma_stages = min(self.subtile_cnt, available_staging_bytes // single_stage_region.nbytes) + self.fc2_staging_spec: Optional[SmemRegion] = self._make_fc2_staging_region(self.fc2_tma_stages) else: self.fc2_tma_stages = 0 self.fc2_staging_spec = None @@ -1218,9 +1082,7 @@ def epilogue_sync_barrier(cls) -> pipeline.NamedBarrier: outside the subtile loop, FC1's amax exchange and FC2's bulk-store handshake belong to different work tiles. """ - return pipeline.NamedBarrier( - barrier_id=cls._EpilogueSyncWaitBarId, num_threads=32 * cls._EpilogueWarpCnt - ) + return pipeline.NamedBarrier(barrier_id=cls._EpilogueSyncWaitBarId, num_threads=32 * cls._EpilogueWarpCnt) def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: """Declare the epilogue scratch as one allocation shared by two exclusive lifetimes. @@ -1275,13 +1137,9 @@ def fc2_tma_staged_smem_spec(self, n_stages: int) -> Tuple[Tuple, Tuple, Tuple[i raise ValueError(f"Unsupported FC2 TMA staging dtype {wire_dtype}.") return shape, stride, swizzle - def fc2_tma_staged_smem_layout( - self, n_stages: int, without_stage_mode: bool = False - ) -> cute.ComposedLayout: + def fc2_tma_staged_smem_layout(self, n_stages: int, without_stage_mode: bool = False) -> cute.ComposedLayout: shape, stride, swizzle = self.fc2_tma_staged_smem_spec(n_stages) - layout = cute.make_composed_layout( - cute.make_swizzle(*swizzle), 0, cute.make_layout(shape, stride=stride) - ) + layout = cute.make_composed_layout(cute.make_swizzle(*swizzle), 0, cute.make_layout(shape, stride=stride)) if without_stage_mode: return cute.select(layout, mode=[0, 1]) return layout @@ -1322,10 +1180,7 @@ def _make_fc2_staging_region(self, stage_count: int) -> SmemRegion: kind="tensor", dtype=single_stage_region.dtype, shape=(*single_stage_region.shape, stage_count), - stride=( - *single_stage_region.stride, - single_stage_region.cosize if stage_count > 1 else 0, - ), + stride=(*single_stage_region.stride, single_stage_region.cosize if stage_count > 1 else 0), swizzle=single_stage_region.swizzle, byte_alignment=single_stage_region.byte_alignment, ) @@ -1399,17 +1254,12 @@ def run( fc2_smem_tensor = smem_workspace.tensor(self.fc2_staging_region, smem_base) else: fc2_smem_tensor = None - if cutlass.const_expr( - self.fc2_use_tma and (tma_atom_fc2_output is None or fc2_tma_output is None) - ): + if cutlass.const_expr(self.fc2_use_tma and (tma_atom_fc2_output is None or fc2_tma_output is None)): raise ValueError("FC2 TMA store requires a TMA atom and token-major output tensor.") if cutlass.const_expr( - self.communication_enabled - and (token_src_metadata is None or peer_rank_ptr_mapper is None) + self.communication_enabled and (token_src_metadata is None or peer_rank_ptr_mapper is None) ): - raise ValueError( - "Communication-enabled Epilogue requires token metadata and a peer pointer mapper." - ) + raise ValueError("Communication-enabled Epilogue requires token metadata and a peer pointer mapper.") if cutlass.const_expr(self.token_back_enabled and fc2_done_counter is None): raise ValueError("Token-back requires an FC2 done counter.") if cutlass.const_expr(self.token_back_push_sf and fc2_output_sf is None): @@ -1488,13 +1338,9 @@ def run( # Publish completion for the work tile snapshotted above. if cur_was_linear1: - flag_tracker = fc1_epi.signal_fc1_done( - prev_work_tile_info, work_tile_info, flag_tracker - ) + flag_tracker = fc1_epi.signal_fc1_done(prev_work_tile_info, work_tile_info, flag_tracker) else: - flag_tracker = fc2_epi.signal_fc2_done( - prev_work_tile_info, work_tile_info, flag_tracker - ) + flag_tracker = fc2_epi.signal_fc2_done(prev_work_tile_info, work_tile_info, flag_tracker) # Tail flush flag_tracker.fire() @@ -1504,9 +1350,7 @@ class _ImmutableAfterInit: def __setattr__(self, name, value): if self.__dict__.get("_frozen_", False): - raise AttributeError( - f"{type(self).__name__} is immutable after __init__ (cannot set {name!r})." - ) + raise AttributeError(f"{type(self).__name__} is immutable after __init__ (cannot set {name!r}).") object.__setattr__(self, name, value) def _freeze(self) -> None: @@ -1535,8 +1379,7 @@ def __init__( # spelled out here, so an 8-bit fc1 output picks up its own atom without a code change. staged_layout = base.fc1_staged_smem_layout(base.subtile_cnt) self.smem_tensor = cute.make_tensor( - cute.recast_ptr(staging_pointer, staged_layout.inner, dtype=base.fc1_output_dtype), - staged_layout.outer, + cute.recast_ptr(staging_pointer, staged_layout.inner, dtype=base.fc1_output_dtype), staged_layout.outer ) # Kept unswizzled and untyped so fc1_quant can carve an fp32 amax-exchange view out of one # staging stage without going through the staging tensor's swizzle. @@ -1572,19 +1415,14 @@ def signal_fc1_done(self, work_tile_info, next_work_tile_info, flag_tracker): or self.intermediate_downproj % self.cluster_tile_intermediate_downproj != 0 ) if cutlass.const_expr(needs_intermediate_guard): - in_bound = ( - work_tile_info.tile_m_idx * self._EpilogueFc1IntermediateDownTileSize - < self.fc1_output.shape[1] - ) + in_bound = work_tile_info.tile_m_idx * self._EpilogueFc1IntermediateDownTileSize < self.fc1_output.shape[1] else: in_bound = True slot = work_tile_info.cumulative_token_block_count + work_tile_info.tile_n_idx flag_address = Int64(0) if in_bound: flag_address = (self.fc1_done_counter.iterator + slot).toint() - return flag_tracker.accumulate( - next_work_tile_info.phase, self.fc1_epi_flag_batch, flag_address - ) + return flag_tracker.accumulate(next_work_tile_info.phase, self.fc1_epi_flag_batch, flag_address) @cute.jit def __call__( @@ -1595,13 +1433,9 @@ def __call__( acc_consumer_state, ): # (tokens_this_expert, intermediate_down, 1) - real_fc1_output, _ = self.kernel_extension.get_gmem_tensor( - "c", self.fc1_output, work_tile_info - ) + real_fc1_output, _ = self.kernel_extension.get_gmem_tensor("c", self.fc1_output, work_tile_info) # (tokens_this_expert, intermediate_down, 1) - real_fc1_output_sf, _ = self.kernel_extension.get_gmem_tensor( - "sfc", self.fc1_output_sf, work_tile_info - ) + real_fc1_output_sf, _ = self.kernel_extension.get_gmem_tensor("sfc", self.fc1_output_sf, work_tile_info) # subtile-irrevalent hoist out here. if cutlass.const_expr(self.optional_epi_args.fc1_alpha is not None): alpha_val = self.optional_epi_args.fc1_alpha[work_tile_info.expert_idx] @@ -1613,8 +1447,7 @@ def __call__( norm_const = None # (cta_tile_m, cta_tile_n) -> (epi_tile_m, epi_tile_n, iters) tmem_acc_tensor_tiled_by_epi_tile = cute.flat_divide( - tmem_acc_tensor, - (self._EpilogueFc1IntermediateGateUpTileSize, self._EpilogueTokenTileSize), + tmem_acc_tensor, (self._EpilogueFc1IntermediateGateUpTileSize, self._EpilogueTokenTileSize) )[None, None, 0, None] acc_pipeline.consumer_wait(acc_consumer_state) @@ -1664,9 +1497,7 @@ def run_subtile( # Step -1: preload topk scores. current_two_token_idices = ( - work_tile_info.tile_n_idx * self.cta_tile_n - + subtile_idx * self._EpilogueTokenTileSize - + self.lane_idx, + work_tile_info.tile_n_idx * self.cta_tile_n + subtile_idx * self._EpilogueTokenTileSize + self.lane_idx, work_tile_info.tile_n_idx * self.cta_tile_n + subtile_idx * self._EpilogueTokenTileSize + self.lane_idx @@ -1687,11 +1518,8 @@ def run_subtile( up_token_32_64 = cute.make_rmem_tensor((16,), cutlass.Float32) # Although hardcode is not right, but since the whole tmem transpose is too tricky, I have to hardcode... # (epi_tile_m, epi_tile_n) -> (warp_local_epi_tile_m, epi_tile_n) - # tmem_subtile_tensor_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[(None, - # self.warp_idx), None] - tmem_subtile_tensor_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[ - (None, 0), None - ] + # tmem_subtile_tensor_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[(None, self.warp_idx), None] + tmem_subtile_tensor_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[(None, 0), None] # (warp_local_epi_tile_m, epi_tile_n) -> (((16, 32), 1), (2, 2)) tmem_subtile_tensor_in_first_load_view = cute.logical_divide( cute.zipped_divide(tmem_subtile_tensor_per_warp, (16, 32)), ((16, 32), 1) @@ -1719,16 +1547,12 @@ def run_subtile( ) # Step 1: perform swiglu on the first part, interleave with the second's 32x32 tmem transpose. - token_0_32_pre_quant_pre_trans = self.alpha_swiglu_clamp( - gate_token_0_32, up_token_0_32, alpha_val - ) + token_0_32_pre_quant_pre_trans = self.alpha_swiglu_clamp(gate_token_0_32, up_token_0_32, alpha_val) # gate_token_32_64 / up_token_32_64 are already in the transpose input # distribution documented by TmemTranspose16x32. token_32_64_tmem_trans = TmemTranspose32x32Inplace( - tmem_subtile_tensor.iterator, - reg_tensor_top=gate_token_32_64, - reg_tensor_bot=up_token_32_64, + tmem_subtile_tensor.iterator, reg_tensor_top=gate_token_32_64, reg_tensor_bot=up_token_32_64 ) # Transpose output: each lane holds (token_1, intermediate_16); tmem_dp @@ -1766,8 +1590,7 @@ def run_subtile( ] # (cta_token, cta_intermediate_down) -> (token_64, intermediate_64) fc1_gmem_subtile_view = cute.flat_divide( - fc1_gmem_cta_view, - (self._EpilogueTokenTileSize, self._EpilogueFc1IntermediateDownTileSize), + fc1_gmem_cta_view, (self._EpilogueTokenTileSize, self._EpilogueFc1IntermediateDownTileSize) )[None, None, subtile_idx, 0] tma_smem_src, tma_gmem_dst = cpasync.tma_partition( self.fc1_tma_atom, @@ -1803,21 +1626,15 @@ def alpha_swiglu_clamp( # raise fires during compilation rather than emitting a runtime branch. for _name, _t in (("gate_rmem", gate_rmem), ("up_rmem", up_rmem)): if cutlass.const_expr(_t.element_type is not cutlass.Float32): - raise TypeError( - f"alpha_swiglu_clamp: {_name} must be Float32, got {_t.element_type}" - ) + raise TypeError(f"alpha_swiglu_clamp: {_name} must be Float32, got {_t.element_type}") if cutlass.const_expr(_t.memspace != AddressSpace.rmem): raise ValueError( f"alpha_swiglu_clamp: {_name} must be a register (rmem) tensor, got address space {_t.memspace}" ) if cutlass.const_expr(cute.rank(_t) != 1): - raise ValueError( - f"alpha_swiglu_clamp: {_name} must be 1D, got rank {cute.rank(_t)}" - ) + raise ValueError(f"alpha_swiglu_clamp: {_name} must be 1D, got rank {cute.rank(_t)}") if cutlass.const_expr(cute.size(_t) % 2 != 0): - raise ValueError( - f"alpha_swiglu_clamp: {_name} element count must be even, got {cute.size(_t)}" - ) + raise ValueError(f"alpha_swiglu_clamp: {_name} element count must be even, got {cute.size(_t)}") if cutlass.const_expr(cute.size(gate_rmem) != cute.size(up_rmem)): raise ValueError( "alpha_swiglu_clamp: gate_rmem and up_rmem must have equal size, got " @@ -1848,10 +1665,7 @@ def alpha_swiglu_clamp( activation_alpha = 1.0 if self.swiglu_alpha is None else self.swiglu_alpha neg_activation_alpha_log2e = -activation_alpha * log2_e - neg_log2e_pair = ( - cutlass.Float32(neg_activation_alpha_log2e), - cutlass.Float32(neg_activation_alpha_log2e), - ) + neg_log2e_pair = (cutlass.Float32(neg_activation_alpha_log2e), cutlass.Float32(neg_activation_alpha_log2e)) one_pair = (cutlass.Float32(1.0), cutlass.Float32(1.0)) if cutlass.const_expr(self.swiglu_beta is not None): beta_pair = (cutlass.Float32(self.swiglu_beta), cutlass.Float32(self.swiglu_beta)) @@ -1930,6 +1744,7 @@ def _sigmoid(p0, p1): sigmoid_pair = _sigmoid(g0, g1) if cutlass.const_expr(self.situ_beta is None): + # SwiGLU: out = (up + swiglu_beta) * gate * sigmoid(swiglu_alpha * gate) # The SwiGLUBias offset follows the up clamp. if cutlass.const_expr(self.swiglu_beta is not None): u0, u1 = cute.arch.add_packed_f32x2((u0, u1), beta_pair) @@ -1975,9 +1790,7 @@ def _sigmoid(p0, p1): def fc1_quant( self, work_tile_info: SwapAbFc12WorkTileInfo, - two_token: Tuple[ - cute.Tensor, cute.Tensor - ], # two rmem tensor, each fp32 @ (token_1, intermediate_16) + two_token: Tuple[cute.Tensor, cute.Tensor], # two rmem tensor, each fp32 @ (token_1, intermediate_16) topk_scores: Optional[Tuple[cutlass.Float32, cutlass.Float32]], norm_const: Optional[cutlass.Float32], intermediate_output_size: cutlass.Int32, @@ -2011,8 +1824,7 @@ def fc1_quant( cute.recast_ptr(self.staging_pointer, dtype=cutlass.Float32) + subtile_idx * cutlass.Int32(self.fc1_staging_stage_bytes // 4), cute.make_layout( - (self.fc1_amax_token_chunks, self._EpilogueWarpCnt * 32), - stride=(1, self.fc1_amax_token_chunks), + (self.fc1_amax_token_chunks, self._EpilogueWarpCnt * 32), stride=(1, self.fc1_amax_token_chunks) ), ), 16, @@ -2039,13 +1851,8 @@ def fc1_quant( stores_scale_factor = self.warp_idx % 2 == 0 else: stores_scale_factor = True - subtile_token_start = ( - work_tile_info.tile_n_idx * self.cta_tile_n + subtile_idx * self._EpilogueTokenTileSize - ) - token_idx_pair = ( - subtile_token_start + self.lane_idx, - subtile_token_start + self.lane_idx + 32, - ) + subtile_token_start = work_tile_info.tile_n_idx * self.cta_tile_n + subtile_idx * self._EpilogueTokenTileSize + token_idx_pair = (subtile_token_start + self.lane_idx, subtile_token_start + self.lane_idx + 32) # This subtile's (token, intermediate) shared output stage, tiled into (1, 16) blocks so # each thread's cells slice out directly (zipped_divide + slice; avoids the ambiguous @@ -2102,9 +1909,7 @@ def fc1_quant( # 4) STS the quantized values into this subtile's shared output stage. # ((1, 16), (token_tile_size, warp_cnt)) -> (16) smem_thread_row = smem_tiled[(0, None), (self.lane_idx + 32 * half, self.warp_idx)] - cute.copy( - store_atom, cute.coalesce(data_by_token[None, half]), cute.coalesce(smem_thread_row) - ) + cute.copy(store_atom, cute.coalesce(data_by_token[None, half]), cute.coalesce(smem_thread_row)) @dataclasses.dataclass(frozen=True) @@ -2199,22 +2004,24 @@ def signal_fc2_done(self, work_tile_info, next_work_tile_info, flag_tracker): publish: cutlass.Constexpr = self.token_back_enabled flag_address = Int64(0) if cutlass.const_expr(publish): - flag_address = (self.fc2_done_counter.iterator + work_tile_info.expert_idx).toint() + if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): # nosec B105 + # The scheduler maps every fallback CTA into the preferred + # logical cluster. N=1 keeps this slot common to all hidden + # shards; padded hidden CTAs must also publish their completion. + slot = work_tile_info.cumulative_token_block_count + work_tile_info.tile_n_idx + flag_address = (self.fc2_done_counter.iterator + slot).toint() + else: + flag_address = (self.fc2_done_counter.iterator + work_tile_info.expert_idx).toint() no_fire: cutlass.Constexpr = not publish - return flag_tracker.accumulate( - next_work_tile_info.phase, self.fc2_epi_flag_batch, flag_address, no_fire - ) + return flag_tracker.accumulate(next_work_tile_info.phase, self.fc2_epi_flag_batch, flag_address, no_fire) @cute.jit def _make_output_router(self, work_tile_info: SwapAbFc12WorkTileInfo) -> "Fc2OutputRouter": task_tile_data_row_start = ( - work_tile_info.cumulative_data_physical_row - + work_tile_info.tile_n_idx * cutlass.Int32(self.cta_tile_n) + work_tile_info.cumulative_data_physical_row + work_tile_info.tile_n_idx * cutlass.Int32(self.cta_tile_n) ) hidden_base_this_cta_tile = work_tile_info.tile_m_idx * cutlass.Int32(self.cta_tile_m) - valid_hidden_this_cta_tile = ( - cutlass.Int32(self.fc2_output.shape[2]) - hidden_base_this_cta_tile - ) + valid_hidden_this_cta_tile = cutlass.Int32(self.fc2_output.shape[2]) - hidden_base_this_cta_tile if valid_hidden_this_cta_tile < 0: valid_hidden_this_cta_tile = 0 if valid_hidden_this_cta_tile > self._EpilogueFc2HiddenTileSize: @@ -2223,9 +2030,7 @@ def _make_output_router(self, work_tile_info: SwapAbFc12WorkTileInfo) -> "Fc2Out metadata = None peer_rank_ptr_mapper = None data_token_base = task_tile_data_row_start - if cutlass.const_expr( - self.token_src_metadata is not None and not self.token_back_push_data - ): + if cutlass.const_expr(self.token_src_metadata is not None and not self.token_back_push_data): metadata = cute.domain_offset((task_tile_data_row_start,), self.token_src_metadata) peer_rank_ptr_mapper = self.peer_rank_ptr_mapper data_token_base = None @@ -2233,10 +2038,7 @@ def _make_output_router(self, work_tile_info: SwapAbFc12WorkTileInfo) -> "Fc2Out if cutlass.const_expr(self.combine_format.is_quantized): base_outputs = (self.fc2_output, self.fc2_output_sf) token_bases = (data_token_base, task_tile_data_row_start) - output_mappings = ( - self.process_pipeline.store_out_mapping, - self.process_pipeline.sf_store_out_mapping, - ) + output_mappings = (self.process_pipeline.store_out_mapping, self.process_pipeline.sf_store_out_mapping) else: base_outputs = self.fc2_output token_bases = data_token_base @@ -2258,11 +2060,7 @@ def _make_output_router(self, work_tile_info: SwapAbFc12WorkTileInfo) -> "Fc2Out @cute.jit def __call__( - self, - work_tile_info: SwapAbFc12WorkTileInfo, - tmem_acc_tensor: cute.Tensor, - acc_pipeline, - acc_consumer_state, + self, work_tile_info: SwapAbFc12WorkTileInfo, tmem_acc_tensor: cute.Tensor, acc_pipeline, acc_consumer_state ): # subtile-irrelevant hoist: fc2 alpha scales raw fc2 accumulators before f2fp. if cutlass.const_expr(self.optional_epi_args.fc2_alpha is not None): @@ -2313,9 +2111,7 @@ def run_subtile( casted = process_pipeline.f2fp(*loaded, alpha_val=alpha_val) # reorder returns a bare RMEM fragment in the store's expected pre-store # distribution; reorder + store are paired 1:1 inside the pipeline. - pre_store = process_pipeline.post_f2fp_reorder( - casted=casted, tmem_subtile_view=tmem_subtile_tensor - ) + pre_store = process_pipeline.post_f2fp_reorder(casted=casted, tmem_subtile_view=tmem_subtile_tensor) process_pipeline.store_function( epi=self, subtile=pre_store, @@ -2344,9 +2140,7 @@ class Fc2OutputRouter: epi_tid: cutlass.Int32 combine_format: CombineFormat # After metadata prefetch - dst_ptrs: Optional[cute.Tensor] = ( - None # i64 x (copy_iters_this_thread_cta_tile), fundamentally the pointers. - ) + dst_ptrs: Optional[cute.Tensor] = None # i64 x (copy_iters_this_thread_cta_tile), fundamentally the pointers. valid: Optional[cute.Tensor] = None # (copy_iters_this_thread_cta_tile) @property @@ -2368,11 +2162,7 @@ def sf_token_base(self) -> Optional[cutlass.Int32]: @property def data_mapping(self) -> FunctionMapping: - return ( - self.output_mappings[0] - if isinstance(self.output_mappings, tuple) - else self.output_mappings - ) + return self.output_mappings[0] if isinstance(self.output_mappings, tuple) else self.output_mappings @property def sf_mapping(self) -> Optional[FunctionMapping]: @@ -2380,9 +2170,7 @@ def sf_mapping(self) -> Optional[FunctionMapping]: def __post_init__(self) -> None: if (self.metadata is None) == (self.data_token_base is None): - raise ValueError( - "Fc2OutputRouter requires exactly one of metadata or a (data) token base." - ) + raise ValueError("Fc2OutputRouter requires exactly one of metadata or a (data) token base.") if (self.metadata is None) != (self.peer_rank_ptr_mapper is None): raise ValueError("Fc2OutputRouter requires peer_rank_ptr_mapper iff metadata is set.") if self.reduce_topk_in_epilogue and self.metadata is None: @@ -2401,8 +2189,7 @@ def prefetch(self) -> "Fc2OutputRouter": valid = cute.make_rmem_tensor((copy_iters,), cutlass.Int32) dst_ptrs = cute.make_rmem_tensor((copy_iters,), cutlass.Int64) - # Compiler should be able to optimize the same token_copy_group's offset add. (Fundamental cse + - # strength_reduce) + # Compiler should be able to optimize the same token_copy_group's offset add. (Fundamental cse + strength_reduce) # We should check the SASS to ensure this happens. for iter_idx in cutlass.range_constexpr(copy_iters): coord = self.data_mapping.evaluate(epi_tid=self.epi_tid, iter_idx=iter_idx) @@ -2421,14 +2208,11 @@ def prefetch(self) -> "Fc2OutputRouter": dst_hidden = hidden_in_tile + self.hidden_base_this_cta_tile # Int64 token coord: dst_tokens*K*H overflows int32 once # T*K*H exceeds 2^31 (data_output is (token, topk, hidden)). - dst_ptrs[iter_idx] = self.data_output[ - Int64(dst_tokens), None, dst_hidden - ].iterator.toint() + dst_ptrs[iter_idx] = self.data_output[Int64(dst_tokens), None, dst_hidden].iterator.toint() else: md = TokenSrcMetadata.load( - self.metadata.iterator.toint() - + Int64(token_in_tile) * Int64(TokenSrcMetadata.nbytes) + self.metadata.iterator.toint() + Int64(token_in_tile) * Int64(TokenSrcMetadata.nbytes) ) dst_rank = md.src_rank dst_token = md.src_token @@ -2440,9 +2224,7 @@ def prefetch(self) -> "Fc2OutputRouter": # Int64 token coord: domain_offset on (token, topk, hidden) # computes dst_token*K*H, which overflows int32 once T*K*H > 2^31. dst_ptrs[iter_idx] = self.peer_rank_ptr_mapper.map_pointer( - cute.domain_offset( - (Int64(dst_token), dst_topk, dst_hidden), self.data_output - ).iterator, + cute.domain_offset((Int64(dst_token), dst_topk, dst_hidden), self.data_output).iterator, dst_rank, byte_alignment=32, ).toint() @@ -2450,9 +2232,7 @@ def prefetch(self) -> "Fc2OutputRouter": return dataclasses.replace(self, dst_ptrs=dst_ptrs, valid=valid) @cute.jit - def get_data_dst( - self, iter_idx: Union[int, cutlass.Int32] - ) -> Tuple[cute.Pointer, cutlass.Int32]: + def get_data_dst(self, iter_idx: Union[int, cutlass.Int32]) -> Tuple[cute.Pointer, cutlass.Int32]: """Per-issue DATA destination: gmem pointer + validity predicate. The router owns ``data_output`` so the caller never re-assembles a @@ -2484,9 +2264,7 @@ def get_data_dst( # comm: read the pointer / validity prefetched by prefetch(). addr = self.dst_ptrs[iter_idx] pred = self.valid[iter_idx] - ptr = cute.make_ptr( - self.data_output.element_type, addr, AddressSpace.gmem, assumed_align=32 - ) + ptr = cute.make_ptr(self.data_output.element_type, addr, AddressSpace.gmem, assumed_align=32) return ptr, pred @cute.jit @@ -2513,9 +2291,7 @@ def get_sf_dst(self, iter_idx: Union[int, cutlass.Int32]) -> Tuple[cute.Pointer, addr = self.sf_output[Int64(sf_row), None, sf_hidden].iterator.toint() # Per-block scale offsets are element-granular; claim the scale dtype's # natural element alignment (e8m0 1 B / bf16 2 B). - sf_ptr = cute.make_ptr( - self.sf_output.element_type, addr, AddressSpace.gmem, assumed_align=4 - ) + sf_ptr = cute.make_ptr(self.sf_output.element_type, addr, AddressSpace.gmem, assumed_align=4) return sf_ptr, pred @@ -2539,16 +2315,13 @@ def make_fc2_stg_cta_store_out_mapping( }, ) store_out_mapping = FunctionMapping( - domain=CoordinateSpace( - ("epi_tid", "iter_idx"), (128, stgs_per_hidden32 * cta_token_tile_size // 32) - ), + domain=CoordinateSpace(("epi_tid", "iter_idx"), (128, stgs_per_hidden32 * cta_token_tile_size // 32)), codomain=CoordinateSpace( ("token_in_cta_tile", "hidden_in_cta_tile"), (cta_token_tile_size, cta_hidden_tile_size) ), function=lambda epi_tid, iter_idx: { "token_in_cta_tile": epi_tid % 32 + iter_idx // stgs_per_hidden32 * 32, - "hidden_in_cta_tile": (iter_idx % stgs_per_hidden32) * elems_per_stg - + epi_tid // 32 * 32, + "hidden_in_cta_tile": (iter_idx % stgs_per_hidden32) * elems_per_stg + epi_tid // 32 * 32, }, ) sf_store_out_mapping = None @@ -2562,8 +2335,7 @@ def stg_sf_mapping(epi_tid, iter_idx): sf_store_out_mapping = FunctionMapping( domain=CoordinateSpace(("epi_tid", "iter_idx"), (128, cta_token_tile_size // 32)), codomain=CoordinateSpace( - ("token_in_cta_tile", "hidden_in_cta_tile"), - (cta_token_tile_size, cta_hidden_tile_size), + ("token_in_cta_tile", "hidden_in_cta_tile"), (cta_token_tile_size, cta_hidden_tile_size) ), function=stg_sf_mapping, ) @@ -2593,10 +2365,7 @@ def make_fc2_redg_cta_store_out_mapping( + (epi_tid % 32) // 4 ), "hidden_in_cta_tile": ( - (epi_tid // 32) * 32 - + (epi_tid % 4) * 4 - + (((elem_idx // 4) % 4) // 2) * 16 - + elem_idx % 4 + (epi_tid // 32) * 32 + (epi_tid % 4) * 4 + (((elem_idx // 4) % 4) // 2) * 16 + elem_idx % 4 ), }, ) @@ -2616,17 +2385,13 @@ def make_fc2_redg_cta_store_out_mapping( + ((iter_idx % 4) % 2) * 8 + (epi_tid % 32) // 4 ), - "hidden_in_cta_tile": ( - (epi_tid // 32) * 32 + (epi_tid % 4) * 4 + ((iter_idx % 4) // 2) * 16 - ), + "hidden_in_cta_tile": ((epi_tid // 32) * 32 + (epi_tid % 4) * 4 + ((iter_idx % 4) // 2) * 16), }, ) return store_out_mapping, None, fundamental_mapping -def make_fc2_ublk_store_out_mapping( - combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int -): +def make_fc2_ublk_store_out_mapping(combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int): assert cta_hidden_tile_size == 128 assert cta_token_tile_size % 64 == 0 assert combine_format.act_dtype.width in (8, 16), "fc2 UBLK wire dtype must be fp8 or bf16." @@ -2652,9 +2417,7 @@ def make_fc2_ublk_store_out_mapping( ("token_in_cta_tile", "hidden_in_cta_tile"), (max_token_cta_tile, cta_hidden_tile_size) ), function=lambda epi_tid, iter_idx: { - "token_in_cta_tile": ( - iter_idx * 128 + ((epi_tid % 32) // 16) * 64 + (epi_tid // 32) * 16 + epi_tid % 16 - ), + "token_in_cta_tile": (iter_idx * 128 + ((epi_tid % 32) // 16) * 64 + (epi_tid // 32) * 16 + epi_tid % 16), "hidden_in_cta_tile": 0, }, ) @@ -2687,9 +2450,7 @@ def fc2_f2fp(*tensors, alpha_val: Optional[cutlass.Float32] = None, **_) -> cute elems_processed = 0 for t in tensors: current_tensor_size = cute.size(t) - dst = cute.make_tensor( - converted_acc.iterator + elems_processed, cute.make_layout((current_tensor_size,)) - ) + dst = cute.make_tensor(converted_acc.iterator + elems_processed, cute.make_layout((current_tensor_size,))) if cutlass.const_expr(alpha_val is None): dst.store(t.load().to(reorder_dtype)) else: @@ -2719,9 +2480,7 @@ def fc2_stg_tmem_acc_load(*, tmem_subtile_tensor: cute.Tensor, **_): ptr = tmem_subtile_tensor.iterator half_lane_offset = 16 * TmemTranspose32x64B16Movm._tmem_row_stride top_view = cute.make_tensor(ptr, TmemTranspose32x64B16Movm._tmem_layout(16, 64)) - bottom_view = cute.make_tensor( - ptr + half_lane_offset, TmemTranspose32x64B16Movm._tmem_layout(16, 64) - ) + bottom_view = cute.make_tensor(ptr + half_lane_offset, TmemTranspose32x64B16Movm._tmem_layout(16, 64)) top = cute.make_rmem_tensor((32,), cutlass.Float32) bottom = cute.make_rmem_tensor((32,), cutlass.Float32) cute.copy(atom_ld16x256, top_view, TmemTranspose32x64B16Movm._rmem_copy_view(top, 32)) @@ -2734,13 +2493,9 @@ def fc2_ublk_tmem_acc_load(*, tmem_subtile_tensor: cute.Tensor, epi, **_): # UBLK consumes a warp-local 32-hidden x 64-token slice. The caller passes # the CTA-level 128-hidden x 64-token subtile view, so select this epi # warp's hidden block before issuing LDTM.x64. - tmem_subtile_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[ - (None, epi.warp_idx), None - ] + tmem_subtile_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[(None, epi.warp_idx), None] raw_regs = cute.make_rmem_tensor((64,), cutlass.Float32) - atom_ld32x32_x64 = cute.make_copy_atom( - tcgen05.Ld32x32bOp(tcgen05.Repetition.x64), cutlass.Float32 - ) + atom_ld32x32_x64 = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x64), cutlass.Float32) cute.copy( atom_ld32x32_x64, wrap_into_copy_standard_layout(tmem_subtile_per_warp), @@ -2765,20 +2520,14 @@ def fc2_stg_post_f2fp_reorder( def fc2_redg_post_f2fp_reorder(*, casted: cute.Tensor, tmem_subtile_view: cute.Tensor, **_): # (epi_tid, elem_idx) -> (token_64, hidden_128), each thread hold token_2 x hidden_32 natural = fc2_stg_post_f2fp_reorder(casted=casted, tmem_subtile_view=tmem_subtile_view) - core_matrix_reorder_sttm_atom = cute.make_copy_atom( - tcgen05.St32x32bOp(tcgen05.Repetition.x16), cutlass.Float32 - ) - core_matrix_reorder_ldtm_atom = cute.make_copy_atom( - tcgen05.Ld16x256bOp(tcgen05.Repetition.x2), cutlass.Float32 - ) + core_matrix_reorder_sttm_atom = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x16), cutlass.Float32) + core_matrix_reorder_ldtm_atom = cute.make_copy_atom(tcgen05.Ld16x256bOp(tcgen05.Repetition.x2), cutlass.Float32) # ((16, 2), token_32_group) token_groups = cute.logical_divide(cute.zipped_divide(natural, (32,)), (16, None)) out = cute.make_rmem_tensor(token_groups.shape, casted.dtype) out_as_i32 = cute.recast_tensor(out, cutlass.Float32) # (32, 64) - tmem_warp = cute.flat_divide(tmem_subtile_view, (32, cute.size(tmem_subtile_view, 1)))[ - None, None, 0, 0 - ] + tmem_warp = cute.flat_divide(tmem_subtile_view, (32, cute.size(tmem_subtile_view, 1)))[None, None, 0, 0] # (16, 16, 16dp_group, token_32_groups). Note, this tmem can provide 2x cols since the original is bf16. tmem_groups = cute.flat_divide(tmem_warp, (16, 16)) for group_idx in cutlass.range_constexpr(cute.size(token_groups, 1)): @@ -2818,9 +2567,7 @@ def fc2_stg_store_function( sf_regs = None stg_width_elems: cutlass.Constexpr[int] = min(32, 256 // data_subtile.element_type.width) stg_bits: cutlass.Constexpr[int] = stg_width_elems * data_subtile.element_type.width - copy_atom_vec = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=stg_bits - ) + copy_atom_vec = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=stg_bits) elems_per_thread: cutlass.Constexpr[int] = cute.size(data_subtile) if cutlass.const_expr(elems_per_thread % stg_width_elems != 0): raise ValueError( @@ -2832,9 +2579,7 @@ def fc2_stg_store_function( sf_scales_per_stg: cutlass.Constexpr[int] = 32 // epi.combine_format.scale_block token_groups_per_subtile: cutlass.Constexpr[int] = epi._EpilogueTokenTileSize // 32 for token_group in cutlass.range_constexpr(token_groups_per_subtile): - sf_iter = ( - cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group - ) + sf_iter = cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group sf_ptr, sf_pred = fc2_output_router.get_sf_dst(sf_iter) if sf_pred != cutlass.Int32(0): sf_dst = cute.make_tensor(sf_ptr, cute.make_layout((sf_scales_per_stg,))) @@ -2851,11 +2596,7 @@ def fc2_stg_store_function( if pred != cutlass.Int32(0): src_i = cute.make_tensor(copy_src[None, local_iter].iterator, single_copy_layout) dst_i = cute.make_tensor(dst_ptr, single_copy_layout) - cute.copy( - copy_atom_vec, - cute.recast_tensor(src_i, cutlass.Int32), - cute.recast_tensor(dst_i, cutlass.Int32), - ) + cute.copy(copy_atom_vec, cute.recast_tensor(src_i, cutlass.Int32), cute.recast_tensor(dst_i, cutlass.Int32)) @cute.jit @@ -2869,12 +2610,8 @@ def fc2_tma_store_function( fc2_output_router: Fc2OutputRouter, **_, ): - if cutlass.const_expr( - epi.smem_tensor is None or epi.tma_atom_fc2_output is None or epi.fc2_tma_output is None - ): - raise ValueError( - "FC2 TMA store requires staged SMEM, a TMA atom, and a token-major output tensor." - ) + if cutlass.const_expr(epi.smem_tensor is None or epi.tma_atom_fc2_output is None or epi.fc2_tma_output is None): + raise ValueError("FC2 TMA store requires staged SMEM, a TMA atom, and a token-major output tensor.") if cutlass.const_expr(epi.combine_format.is_quantized): data_subtile, sf_regs = QuantImpl(epi.combine_format, "regs_in_thread")(subtile) @@ -2886,9 +2623,7 @@ def fc2_tma_store_function( sf_scales_per_stg: cutlass.Constexpr[int] = 32 // epi.combine_format.scale_block token_groups_per_subtile: cutlass.Constexpr[int] = epi._EpilogueTokenTileSize // 32 for token_group in cutlass.range_constexpr(token_groups_per_subtile): - sf_iter = ( - cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group - ) + sf_iter = cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group sf_ptr, sf_pred = fc2_output_router.get_sf_dst(sf_iter) if sf_pred != cutlass.Int32(0): sf_dst = cute.make_tensor(sf_ptr, cute.make_layout((sf_scales_per_stg,))) @@ -2902,15 +2637,11 @@ def fc2_tma_store_function( or data_subtile.stride[0] != 1 or 32 % vector_elements != 0 ): - raise ValueError( - "FC2 TMA store requires a contiguous two-token by 32-hidden register fragment." - ) + raise ValueError("FC2 TMA store requires a contiguous two-token by 32-hidden register fragment.") stage_idx = epilogue_iter_idx % cutlass.Int32(epi.fc2_tma_stages) smem_stage = epi.smem_tensor[None, None, stage_idx] - store_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), data_subtile.element_type, num_bits_per_copy=128 - ) + store_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), data_subtile.element_type, num_bits_per_copy=128) tiler_mn = (64, 128) layout_copy_tv = cute.make_layout(((32, 4), (32, 2)), stride=((1, 2048), (64, 32))) tiled_store = cute.make_tiled_copy(store_atom, layout_copy_tv, tiler_mn) @@ -2924,9 +2655,7 @@ def fc2_tma_store_function( + work_tile_info.tile_n_idx * cutlass.Int32(epi.cta_tile_n) + subtile_idx * cutlass.Int32(epi._EpilogueTokenTileSize) ) // cutlass.Int32(epi._EpilogueTokenTileSize) - tiled_output = cute.flat_divide( - epi.fc2_tma_output, (epi._EpilogueTokenTileSize, epi._EpilogueFc2HiddenTileSize) - ) + tiled_output = cute.flat_divide(epi.fc2_tma_output, (epi._EpilogueTokenTileSize, epi._EpilogueFc2HiddenTileSize)) gmem_subtile = tiled_output[None, None, token_tile_idx, work_tile_info.tile_m_idx, 0] tma_smem_src, tma_gmem_dst = cpasync.tma_partition( epi.tma_atom_fc2_output, @@ -2965,9 +2694,7 @@ def fc2_ublk_store_function_impl( sf_scales_per_stg: cutlass.Constexpr[int] = 32 // epi.combine_format.scale_block token_groups_per_subtile: cutlass.Constexpr[int] = epi._EpilogueTokenTileSize // 32 for token_group in cutlass.range_constexpr(token_groups_per_subtile): - sf_iter = ( - cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group - ) + sf_iter = cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group sf_ptr, sf_pred = fc2_output_router.get_sf_dst(sf_iter) if sf_pred != cutlass.Int32(0): sf_dst = cute.make_tensor(sf_ptr, cute.make_layout((sf_scales_per_stg,))) @@ -2990,12 +2717,8 @@ def fc2_ublk_store_function_impl( or data_subtile.stride[0] != 1 or 32 % vector_elements != 0 ): - raise ValueError( - "FC2 UBLK store requires a contiguous two-token by 32-hidden register fragment." - ) - store_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), data_subtile.element_type, num_bits_per_copy=128 - ) + raise ValueError("FC2 UBLK store requires a contiguous two-token by 32-hidden register fragment.") + store_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), data_subtile.element_type, num_bits_per_copy=128) tiler_mn = (64, 128) layout_copy_tv = cute.make_layout(((32, 4), (32, 2)), stride=((1, 2048), (64, 32))) tiled_store = cute.make_tiled_copy(store_atom, layout_copy_tv, tiler_mn) @@ -3005,9 +2728,7 @@ def fc2_ublk_store_function_impl( cute.copy(tiled_store, register_partition, smem_partition) else: if cutlass.const_expr(cute.size(data_subtile) != epi._EpilogueTokenTileSize): - raise ValueError( - "BF16 UBLK staging requires one register per token in the 64-token subtile." - ) + raise ValueError("BF16 UBLK staging requires one register per token in the 64-token subtile.") warp_hidden_base = cutlass.Int32(warp_idx * 32) for token_idx in cutlass.range_constexpr(epi._EpilogueTokenTileSize): smem_stage[token_idx, warp_hidden_base + lane_idx] = data_subtile[token_idx] @@ -3063,18 +2784,14 @@ def fc2_redg_store_function( if pred != cutlass.Int32(0): bf16x4 = subtile_by_redg_issue[None, local_iter] packed_bf16x2 = cute.recast_tensor(bf16x4, cutlass.Float32) - red_add_relaxed_sys_v2_bf16x2( - dst_ptr, cutlass.Float32(packed_bf16x2[0]), cutlass.Float32(packed_bf16x2[1]) - ) + red_add_relaxed_sys_v2_bf16x2(dst_ptr, cutlass.Float32(packed_bf16x2[0]), cutlass.Float32(packed_bf16x2[1])) def make_fc2_stg_process_pipeline( *, combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int ) -> Fc2ProcessPipeline: - store_out_mapping, sf_store_out_mapping, fundamental_mapping = ( - make_fc2_stg_cta_store_out_mapping( - combine_format, cta_token_tile_size, cta_hidden_tile_size - ) + store_out_mapping, sf_store_out_mapping, fundamental_mapping = make_fc2_stg_cta_store_out_mapping( + combine_format, cta_token_tile_size, cta_hidden_tile_size ) return Fc2ProcessPipeline( tmem_acc_load=fc2_stg_tmem_acc_load, @@ -3090,10 +2807,8 @@ def make_fc2_stg_process_pipeline( def make_fc2_tma_process_pipeline( *, combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int ) -> Fc2ProcessPipeline: - store_out_mapping, sf_store_out_mapping, fundamental_mapping = ( - make_fc2_stg_cta_store_out_mapping( - combine_format, cta_token_tile_size, cta_hidden_tile_size - ) + store_out_mapping, sf_store_out_mapping, fundamental_mapping = make_fc2_stg_cta_store_out_mapping( + combine_format, cta_token_tile_size, cta_hidden_tile_size ) return Fc2ProcessPipeline( tmem_acc_load=fc2_stg_tmem_acc_load, @@ -3109,10 +2824,8 @@ def make_fc2_tma_process_pipeline( def make_fc2_redg_process_pipeline( *, combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int ) -> Fc2ProcessPipeline: - store_out_mapping, sf_store_out_mapping, fundamental_mapping = ( - make_fc2_redg_cta_store_out_mapping( - combine_format, cta_token_tile_size, cta_hidden_tile_size - ) + store_out_mapping, sf_store_out_mapping, fundamental_mapping = make_fc2_redg_cta_store_out_mapping( + combine_format, cta_token_tile_size, cta_hidden_tile_size ) return Fc2ProcessPipeline( tmem_acc_load=fc2_stg_tmem_acc_load, diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_extension.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_extension.py index c17849379ec3..a8900c010327 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_extension.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_extension.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Rubin block-scaled SwapAB adapter between FC12 scheduling and kernel tensor views.""" import dataclasses @@ -16,6 +13,7 @@ from .....helpers.dsl_helpers import spin_peek, spin_wait from ....schedulers.fc12_mapping import BlockPhase, SwapAbFc12WorkTileInfo, peek_ready_bit + TensorRole = Literal["a", "b", "sfa", "sfb", "c", "sfc", "topk"] @@ -59,18 +57,14 @@ def __new_from_mlir_values__(self, values: List[ir.Value]) -> "BlockScaledSwapAb def rebuild(field): nonlocal value_index field_value_count = len(extract_mlir_values(field)) - result = new_from_mlir_values( - field, values[value_index : value_index + field_value_count] - ) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) value_index += field_value_count return result fc1_done_counter_pointer = rebuild(self.fc1_done_counter_pointer) fc2_spin_threshold = rebuild(self.fc2_spin_threshold) fc1_ready_counter_pointer = ( - rebuild(self.fc1_ready_counter_pointer) - if self.fc1_ready_counter_pointer is not None - else None + rebuild(self.fc1_ready_counter_pointer) if self.fc1_ready_counter_pointer is not None else None ) expert_remap = rebuild(self.expert_remap) if self.expert_remap is not None else None if value_index != len(values): @@ -110,9 +104,7 @@ def prepare_work_tile(self, work_tile: SwapAbFc12WorkTileInfo) -> SwapAbFc12Work if is_fc1: counter_pointer = self.fc1_ready_counter_pointer + counter_slot peek_flag = Int32(0) - if spin_peek( - counter_pointer, lambda value: value >= work_tile.valid_tokens_in_cta_tile - ): + if spin_peek(counter_pointer, lambda value: value >= work_tile.valid_tokens_in_cta_tile): peek_flag = Int32(peek_ready_bit) phase_and_flags = work_tile.phase_and_flags | peek_flag @@ -170,19 +162,13 @@ def get_gmem_tensor( result = cute.domain_offset((0, 0, expert_idx), tensor) per_expert_shape = (shape[0], shape[1], singleton) sf_layout = tile_atom_to_shape_SF(per_expert_shape, self.sf_vec_size) - return ( - cute.make_tensor(result.iterator, cute.make_layout(sf_layout.shape, stride=stride)), - None, - ) + return (cute.make_tensor(result.iterator, cute.make_layout(sf_layout.shape, stride=stride)), None) if cutlass.const_expr(tensor_name in ("sfb", "sfc")): result = cute.domain_offset((sf_token_offset, 0, 0), tensor) per_expert_shape = (shape[0], shape[1], singleton) sf_layout = tile_atom_to_shape_SF(per_expert_shape, self.sf_vec_size) - return ( - cute.make_tensor(result.iterator, cute.make_layout(sf_layout.shape, stride=stride)), - None, - ) + return (cute.make_tensor(result.iterator, cute.make_layout(sf_layout.shape, stride=stride)), None) if cutlass.const_expr(tensor_name == "c"): result = cute.domain_offset((data_token_offset, 0, 0), tensor) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_mainloop.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_mainloop.py index b26743df47e8..bcd101a404fe 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_mainloop.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_mainloop.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Mainloop component for the Rubin block-scaled swap-AB FC12 kernel.""" from typing import ClassVar, Optional, Tuple @@ -15,13 +12,7 @@ from cutlass.cute.nvgpu import OperandMajorMode, cpasync, tcgen05 from cutlass.cutlass_dsl import Boolean, Int32 -from .....api import ( - ImplDesc, - KernelComponent, - OptionalRequirement, - ProblemDesc, - StaticOrRuntimeIntegerType, -) +from .....api import ImplDesc, KernelComponent, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType from .....helpers.cute_py_helpers import ( Tcgen05MmaInstruction, make_smem_layouts, @@ -33,15 +24,15 @@ from .....helpers.smem_workspace import SmemWorkspace from .....helpers.utils import ceil_div, round_up, strides_equal_ignoring_singletons from .....quant_def import QuantKind -from ....blackwell.custom_mix_cga_helpers import ( +from ....schedulers.base import SchedulerConsumer +from ....schedulers.fc12_mapping import BlockPhase +from ...custom_mix_cga_helpers import ( PipelineTmaUmmaMixedCga, TmaAtomOrPair, bind_executable_tma_load_fields, make_executable_tma_atom_or_pair, select_tma_atom_for_cluster, ) -from ....schedulers.base import SchedulerConsumer -from ....schedulers.fc12_mapping import BlockPhase from . import dynamic_mainloop from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension @@ -49,20 +40,14 @@ class BlockScaledSwapAbFc12Mainloop(KernelComponent): """Own all Rubin load, MMA, pipeline, SMEM, and TMEM mainloop state.""" - ab_pipeline_mbarriers_region: ClassVar[str] = ( - "rubin.swap_ab_fc12.mainloop.ab_pipeline_mbarriers" - ) + ab_pipeline_mbarriers_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.ab_pipeline_mbarriers" a_smem_tensor_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.a_smem_tensor" b_smem_tensor_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.b_smem_tensor" sfa_smem_tensor_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.sfa_smem_tensor" sfb_smem_tensor_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.sfb_smem_tensor" - acc_pipeline_mbarriers_region: ClassVar[str] = ( - "rubin.swap_ab_fc12.mainloop.acc_pipeline_mbarriers" - ) + acc_pipeline_mbarriers_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.acc_pipeline_mbarriers" tmem_holding_buffer_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.tmem_holding_buffer" - tmem_deallocation_mbarrier_region: ClassVar[str] = ( - "rubin.swap_ab_fc12.mainloop.tmem_deallocation_mbarrier" - ) + tmem_deallocation_mbarrier_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.tmem_deallocation_mbarrier" tmem_allocation_barrier_id: ClassVar[int] = 2 @classmethod @@ -108,9 +93,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.fallback_cluster_shape_mn = impl_desc.get("fallback_cluster_shape_mn") self.use_2cta_instrs = impl_desc["use_2cta_instrs"] self.mainloop_smem_budget_bytes = impl_desc["mainloop_smem_budget_bytes"] - self.num_accumulator_consumer_warps_per_cta = impl_desc[ - "num_accumulator_consumer_warps_per_cta" - ] + self.num_accumulator_consumer_warps_per_cta = impl_desc["num_accumulator_consumer_warps_per_cta"] self.architecture = "sm_107" self.mma_cta_count = 2 if self.use_2cta_instrs else 1 @@ -131,15 +114,9 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: if tmem_plan.accumulator_stage_stride_columns != tmem_plan.accumulator_stage_columns: raise ValueError("Rubin MegaMoE does not support overlapping accumulator stages.") if tmem_plan.accumulator_pipeline_stages != tmem_plan.accumulator_stage_count: - raise ValueError( - "Rubin MegaMoE requires one pipeline stage per disjoint accumulator stage." - ) + raise ValueError("Rubin MegaMoE requires one pipeline stage per disjoint accumulator stage.") self.num_accumulator_tmem_cols = tmem_plan.accumulator_columns - self.accumulator_shape = ( - self.cta_tile_m, - self.cta_tile_n, - tmem_plan.accumulator_stage_count, - ) + self.accumulator_shape = (self.cta_tile_m, self.cta_tile_n, tmem_plan.accumulator_stage_count) self.accumulator_stride = (1 << 16, 1, tmem_plan.accumulator_stage_stride_columns) self.mma_instruction = Tcgen05MmaInstruction( @@ -152,12 +129,8 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: sfb_type=self.sf_dtype, sf_vec_size=self.sf_vec_size, ) - self.a_smem_alloc_dtype = tcgen05_smem_alloc_type( - self.a_dtype, self.b_dtype, self.architecture - ) - self.b_smem_alloc_dtype = tcgen05_smem_alloc_type( - self.b_dtype, self.a_dtype, self.architecture - ) + self.a_smem_alloc_dtype = tcgen05_smem_alloc_type(self.a_dtype, self.b_dtype, self.architecture) + self.b_smem_alloc_dtype = tcgen05_smem_alloc_type(self.b_dtype, self.a_dtype, self.architecture) self.sfb_instruction_shape_mnk = ( self.mma_instruction_mnk[0] // self.mma_cta_count, round_up(self.mma_instruction_mnk[1], 128), @@ -173,11 +146,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: sfb_type=self.sf_dtype, sf_vec_size=self.sf_vec_size, ) - self.mma_tiler_sfb = ( - self.sfb_instruction_shape_mnk[0], - self.sfb_instruction_shape_mnk[1], - self.mma_tile_k, - ) + self.mma_tiler_sfb = (self.sfb_instruction_shape_mnk[0], self.sfb_instruction_shape_mnk[1], self.mma_tile_k) self.cluster_layout_shape_vmnk = ( (self.mma_cta_count,), self.cluster_shape_mn[0] // self.mma_cta_count, @@ -186,9 +155,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: ) self.cluster_layout_sfb_shape_vmnk = ((1,), *self.cluster_shape_mn, 1) self.resolved_fallback_cluster_shape_mn = ( - self.cluster_shape_mn - if self.fallback_cluster_shape_mn is None - else self.fallback_cluster_shape_mn + self.cluster_shape_mn if self.fallback_cluster_shape_mn is None else self.fallback_cluster_shape_mn ) self.fallback_cluster_layout_shape_vmnk = ( (self.mma_cta_count,), @@ -196,11 +163,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.resolved_fallback_cluster_shape_mn[1], 1, ) - self.fallback_cluster_layout_sfb_shape_vmnk = ( - (1,), - *self.resolved_fallback_cluster_shape_mn, - 1, - ) + self.fallback_cluster_layout_sfb_shape_vmnk = ((1,), *self.resolved_fallback_cluster_shape_mn, 1) self.is_mixed_cga = self.resolved_fallback_cluster_shape_mn != self.cluster_shape_mn self.num_mcast_ctas_a = self.cluster_layout_shape_vmnk[2] self.num_mcast_ctas_b = self.cluster_layout_shape_vmnk[1] @@ -212,9 +175,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: def _validate_configuration(self) -> None: if self.mma_k_mode != "2x": - raise ValueError( - f"Rubin MegaMoE only supports mma_k_mode='2x', got {self.mma_k_mode!r}." - ) + raise ValueError(f"Rubin MegaMoE only supports mma_k_mode='2x', got {self.mma_k_mode!r}.") for field_name, dimensions, expected_rank in ( ("mma_instruction_mnk", self.mma_instruction_mnk, 3), ("mma_tiler_mnk", self.mma_tiler_mnk, 3), @@ -222,10 +183,7 @@ def _validate_configuration(self) -> None: ): if len(dimensions) != expected_rank: raise ValueError(f"{field_name} must contain {expected_rank} dimensions.") - if not all( - isinstance(dimension, int) and not isinstance(dimension, bool) - for dimension in dimensions - ): + if not all(isinstance(dimension, int) and not isinstance(dimension, bool) for dimension in dimensions): raise TypeError(f"{field_name} dimensions must be Python integers.") if any(dimension <= 0 for dimension in dimensions): raise ValueError(f"{field_name} dimensions must be positive.") @@ -242,9 +200,7 @@ def _validate_configuration(self) -> None: f"got {self.mma_instruction_mnk[0]}." ) if self.mma_instruction_mnk[:2] != self.mma_tiler_mnk[:2]: - raise NotImplementedError( - "Rubin MegaMoE does not implement M/N instruction repetition or B-reuse." - ) + raise NotImplementedError("Rubin MegaMoE does not implement M/N instruction repetition or B-reuse.") if self.mma_tiler_mnk[1] not in (64, 128, 256): raise ValueError("Rubin MegaMoE supports tile N in (64, 128, 256).") if self.mma_tiler_mnk[2] % self.instruction_k != 0: @@ -269,9 +225,7 @@ def _validate_configuration(self) -> None: self.cluster_shape_mn, self.fallback_cluster_shape_mn ) ): - raise ValueError( - "Preferred cluster dimensions must be divisible by fallback dimensions." - ) + raise ValueError("Preferred cluster dimensions must be divisible by fallback dimensions.") if self.mainloop_smem_budget_bytes <= 0: raise ValueError("mainloop_smem_budget_bytes must be positive.") if self.num_accumulator_consumer_warps_per_cta <= 0: @@ -286,9 +240,7 @@ def _validate_configuration(self) -> None: def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: """Register the selected Rubin mainloop SMEM plan.""" - smem_workspace.register_mbarrier( - self.ab_pipeline_mbarriers_region, self.num_ab_pipeline_stages * 2 - ) + smem_workspace.register_mbarrier(self.ab_pipeline_mbarriers_region, self.num_ab_pipeline_stages * 2) for region_name, region in ( (self.a_smem_tensor_region, self.a_smem_region), (self.b_smem_tensor_region, self.b_smem_region), @@ -303,37 +255,23 @@ def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: swizzle=region.swizzle, byte_alignment=region.byte_alignment, ) - smem_workspace.register_mbarrier( - self.acc_pipeline_mbarriers_region, self.num_accumulator_pipeline_stages * 2 - ) + smem_workspace.register_mbarrier(self.acc_pipeline_mbarriers_region, self.num_accumulator_pipeline_stages * 2) smem_workspace.register_mbarrier(self.tmem_deallocation_mbarrier_region, 1) - smem_workspace.register_tensor( - self.tmem_holding_buffer_region, cutlass.Int32, (1,), byte_alignment=8 - ) + smem_workspace.register_tensor(self.tmem_holding_buffer_region, cutlass.Int32, (1,), byte_alignment=8) def _select_ab_stage_plan(self) -> None: stage_one_regions = make_smem_layouts( - self.mma_instruction, - self.mma_tiler_mnk, - 1, - (self.a_major_mode, self.b_major_mode), - self.architecture, + self.mma_instruction, self.mma_tiler_mnk, 1, (self.a_major_mode, self.b_major_mode), self.architecture ) independent_sfb_region = make_smem_layouts( - self.sfb_mma_instruction, - self.mma_tiler_sfb, - 1, - (self.a_major_mode, self.b_major_mode), - self.architecture, + self.sfb_mma_instruction, self.mma_tiler_sfb, 1, (self.a_major_mode, self.b_major_mode), self.architecture )[3] if ( independent_sfb_region.shape != stage_one_regions[3].shape or independent_sfb_region.stride != stage_one_regions[3].stride or independent_sfb_region.nbytes != stage_one_regions[3].nbytes ): - raise ValueError( - "The Rubin independent SFB plan does not match the primary MMA SFB operand." - ) + raise ValueError("The Rubin independent SFB plan does not match the primary MMA SFB operand.") def packed_bytes(region, dtype) -> int: return (region.cosize * int(dtype.width) + 7) // 8 @@ -350,31 +288,22 @@ def packed_bytes(region, dtype) -> int: mbarrier_bytes = int(cutlass.Int64.width) // 8 ab_stage_cost_bytes = self.ab_stage_payload_bytes + 2 * mbarrier_bytes plan_tail_bytes = (2 * self.num_accumulator_pipeline_stages + 2) * mbarrier_bytes - self.num_ab_pipeline_stages = ( - self.mainloop_smem_budget_bytes - plan_tail_bytes - ) // ab_stage_cost_bytes + self.num_ab_pipeline_stages = (self.mainloop_smem_budget_bytes - plan_tail_bytes) // ab_stage_cost_bytes if self.num_ab_pipeline_stages < 1: raise ValueError( f"One AB stage needs {ab_stage_cost_bytes + plan_tail_bytes} bytes, exceeding the " f"{self.mainloop_smem_budget_bytes}-byte mainloop budget." ) - (self.a_smem_region, self.b_smem_region, self.sfa_smem_region, self.sfb_smem_region) = ( - make_smem_layouts( - self.mma_instruction, - self.mma_tiler_mnk, - self.num_ab_pipeline_stages, - (self.a_major_mode, self.b_major_mode), - self.architecture, - ) + (self.a_smem_region, self.b_smem_region, self.sfa_smem_region, self.sfb_smem_region) = make_smem_layouts( + self.mma_instruction, + self.mma_tiler_mnk, + self.num_ab_pipeline_stages, + (self.a_major_mode, self.b_major_mode), + self.architecture, ) self.selected_smem_bytes = sum( region.nbytes - for region in ( - self.a_smem_region, - self.b_smem_region, - self.sfa_smem_region, - self.sfb_smem_region, - ) + for region in (self.a_smem_region, self.b_smem_region, self.sfa_smem_region, self.sfb_smem_region) ) def make_tiled_mma(self) -> cute.TiledMma: @@ -412,12 +341,10 @@ def materialize_codegen_members(self) -> None: cute.make_layout((*self.cluster_shape_mn, 1)), (self.tiled_mma_sfb.thr_id.shape,) ) self.fallback_cluster_layout_vmnk = cute.tiled_divide( - cute.make_layout((*self.resolved_fallback_cluster_shape_mn, 1)), - (self.tiled_mma.thr_id.shape,), + cute.make_layout((*self.resolved_fallback_cluster_shape_mn, 1)), (self.tiled_mma.thr_id.shape,) ) self.fallback_cluster_layout_sfb_vmnk = cute.tiled_divide( - cute.make_layout((*self.resolved_fallback_cluster_shape_mn, 1)), - (self.tiled_mma_sfb.thr_id.shape,), + cute.make_layout((*self.resolved_fallback_cluster_shape_mn, 1)), (self.tiled_mma_sfb.thr_id.shape,) ) if self.cluster_layout_vmnk.shape != self.cluster_layout_shape_vmnk: raise ValueError( @@ -434,10 +361,7 @@ def materialize_codegen_members(self) -> None: f"Fallback cluster layout mismatch: Python {self.fallback_cluster_layout_shape_vmnk}, " f"CuTe {self.fallback_cluster_layout_vmnk.shape}." ) - if ( - self.fallback_cluster_layout_sfb_vmnk.shape - != self.fallback_cluster_layout_sfb_shape_vmnk - ): + if self.fallback_cluster_layout_sfb_vmnk.shape != self.fallback_cluster_layout_sfb_shape_vmnk: raise ValueError( f"Fallback SFB cluster layout mismatch: Python {self.fallback_cluster_layout_sfb_shape_vmnk}, " f"CuTe {self.fallback_cluster_layout_sfb_vmnk.shape}." @@ -463,11 +387,8 @@ def materialize_codegen_members(self) -> None: ("SFA", self.sfa_smem_layout_staged, self.sfa_smem_region), ("SFB", self.sfb_smem_layout_staged, self.sfb_smem_region), ): - if ( - materialized_layout.shape != planned_region.shape - or not strides_equal_ignoring_singletons( - planned_region.shape, materialized_layout.stride, planned_region.stride - ) + if materialized_layout.shape != planned_region.shape or not strides_equal_ignoring_singletons( + planned_region.shape, materialized_layout.stride, planned_region.stride ): raise ValueError( f"{tensor_name} SMEM plan mismatch: Python " @@ -493,24 +414,12 @@ def prepare_tma_load_params( sfa_stage_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0)) sfb_stage_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0)) - a_operation = sm100_utils.cluster_shape_to_tma_atom_A( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - b_operation = sm100_utils.cluster_shape_to_tma_atom_B( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - sfa_operation = sm100_utils.cluster_shape_to_tma_atom_A( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - sfb_operation = sm100_utils.cluster_shape_to_tma_atom_SFB( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - a_internal_type = ( - self.a_smem_alloc_dtype if self.a_smem_alloc_dtype is not self.a_dtype else None - ) - b_internal_type = ( - self.b_smem_alloc_dtype if self.b_smem_alloc_dtype is not self.b_dtype else None - ) + a_operation = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, self.tiled_mma.thr_id) + b_operation = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, self.tiled_mma.thr_id) + sfa_operation = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, self.tiled_mma.thr_id) + sfb_operation = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, self.tiled_mma.thr_id) + a_internal_type = self.a_smem_alloc_dtype if self.a_smem_alloc_dtype is not self.a_dtype else None + b_internal_type = self.b_smem_alloc_dtype if self.b_smem_alloc_dtype is not self.b_dtype else None def make_tma_maybe_pair( make_atom, @@ -678,17 +587,10 @@ def assign_device_members( ) self.fc1_k_tile_count = ceil_div(bound_hidden_size, self.mma_tile_k) self.fc2_k_tile_count = ceil_div(bound_intermediate_gateup_size // 2, self.mma_tile_k) - self.sfb_vmnk_coord = ( - 0, - cta_coord_in_cluster[0], - cta_coord_in_cluster[1], - cta_coord_in_cluster[2], - ) + self.sfb_vmnk_coord = (0, cta_coord_in_cluster[0], cta_coord_in_cluster[1], cta_coord_in_cluster[2]) @cute.jit - def create_ab_pipeline( - self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer - ) -> pipeline.PipelineAsync: + def create_ab_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> pipeline.PipelineAsync: """Create the shared two-producer TMA-to-UMMA pipeline.""" producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 2) return PipelineTmaUmmaMixedCga.create( @@ -709,9 +611,7 @@ def create_ab_pipeline( def create_acc_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer): """Create the UMMA-to-epilogue accumulator pipeline.""" producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) - consumer_thread_count = ( - self.num_accumulator_consumer_warps_per_cta * 32 * self.mma_cta_count - ) + consumer_thread_count = self.num_accumulator_consumer_warps_per_cta * 32 * self.mma_cta_count consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, consumer_thread_count) return pipeline.PipelineUmmaAsync.create( barrier_storage=smem_workspace.ptr(self.acc_pipeline_mbarriers_region, smem_base), @@ -723,9 +623,7 @@ def create_acc_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Poi ) @cute.jit - def create_tmem_allocator( - self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer, *, allocator_warp_id: int - ): + def create_tmem_allocator(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer, *, allocator_warp_id: int): """Bind the SM107 TMEM allocator to finalized SMEM regions.""" allocation_barrier = pipeline.NamedBarrier( barrier_id=self.tmem_allocation_barrier_id, @@ -736,9 +634,7 @@ def create_tmem_allocator( barrier_for_retrieve=allocation_barrier, allocator_warp_id=allocator_warp_id, is_two_cta=self.use_2cta_instrs, - two_cta_tmem_dealloc_mbar_ptr=smem_workspace.ptr( - self.tmem_deallocation_mbarrier_region, smem_base - ), + two_cta_tmem_dealloc_mbar_ptr=smem_workspace.ptr(self.tmem_deallocation_mbarrier_region, smem_base), arch=self.architecture, ) @@ -752,18 +648,12 @@ def _s2t_copy_and_partition(self, smem_tensor: cute.Tensor, tmem_tensor: cute.Te mn_mode = cute.get(compact_smem_tensor.layout, mode=[0, 0]) mn_mode = cute.append(mn_mode, cute.make_layout((4,), stride=(0,))) - broadcast_layout = cute.append( - cute.group_modes(mn_mode, 0), cute.get(compact_smem_tensor.layout, mode=[0, 1]) - ) + broadcast_layout = cute.append(cute.group_modes(mn_mode, 0), cute.get(compact_smem_tensor.layout, mode=[0, 1])) broadcast_layout = cute.append( cute.group_modes(broadcast_layout, 0), cute.get(compact_smem_tensor.layout, mode=[1]) ) - broadcast_layout = cute.append( - broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[2]) - ) - broadcast_layout = cute.append( - broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[3]) - ) + broadcast_layout = cute.append(broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[2])) + broadcast_layout = cute.append(broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[3])) broadcast_smem_tensor = cute.make_tensor(compact_smem_tensor.iterator, broadcast_layout) partitioned_smem = thread_copy.partition_S(broadcast_smem_tensor) @@ -792,17 +682,10 @@ def run_tma_a( multicast_mask = None if cutlass.const_expr(self.is_a_mcast or self.use_2cta_instrs): multicast_mask = tma_multicast_mask( - self.cluster_shape_mn, - None, - self.cta_coord_in_cluster, - None, - self.use_2cta_instrs, - "a", + self.cluster_shape_mn, None, self.cta_coord_in_cluster, None, self.use_2cta_instrs, "a" ) - a_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_vmnk, (0, 0, None, 0)).shape - ) + a_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_vmnk, (0, 0, None, 0)).shape) sfa_cta_layout = a_cta_layout thread_mma = self.tiled_mma.get_slice(self.mma_cta_index) @@ -812,16 +695,10 @@ def run_tma_a( if is_fc1: iket.range_push("tma_weight_fc1") k_tile_count = self.fc1_k_tile_count - real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor( - "a", fc1_tma_a_tensor, work_tile - ) - real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfa", fc1_tma_sfa_tensor, work_tile - ) + real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor("a", fc1_tma_a_tensor, work_tile) + real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor("sfa", fc1_tma_sfa_tensor, work_tile) - global_a = cute.local_tile( - real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) - ) + global_a = cute.local_tile(real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None)) global_sfa = cute.local_tile( real_sfa, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) ) @@ -876,16 +753,10 @@ def run_tma_a( else: iket.range_push("tma_weight_fc2") k_tile_count = self.fc2_k_tile_count - real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor( - "a", fc2_tma_a_tensor, work_tile - ) - real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfa", fc2_tma_sfa_tensor, work_tile - ) + real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor("a", fc2_tma_a_tensor, work_tile) + real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor("sfa", fc2_tma_sfa_tensor, work_tile) - global_a = cute.local_tile( - real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) - ) + global_a = cute.local_tile(real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None)) global_sfa = cute.local_tile( real_sfa, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) ) @@ -1006,19 +877,13 @@ def run_tma_b( "sfb", ) - b_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_vmnk, (0, None, 0, 0)).shape - ) - sfb_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape - ) + b_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_vmnk, (0, None, 0, 0)).shape) + sfb_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape) b_partition_coord = self.main_vmnk_coord[1] sfb_partition_coord = self.sfb_vmnk_coord[1] if cutlass.const_expr(self.is_mixed_cga): if self.is_fallback_cluster: - split_factor = ( - self.cluster_shape_mn[0] // self.resolved_fallback_cluster_shape_mn[0] - ) + split_factor = self.cluster_shape_mn[0] // self.resolved_fallback_cluster_shape_mn[0] b_partition_coord = b_partition_coord * Int32(split_factor) sfb_partition_coord = sfb_partition_coord * Int32(split_factor) thread_mma = self.tiled_mma.get_slice(self.mma_cta_index) @@ -1033,24 +898,17 @@ def run_tma_b( kernel_extension.wait_for_input(work_tile) iket.range_pop() k_tile_count = self.fc1_k_tile_count - real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor( - "b", fc1_tma_b_tensor, work_tile - ) - real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfb", fc1_tma_sfb_tensor, work_tile - ) + real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor("b", fc1_tma_b_tensor, work_tile) + real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor("sfb", fc1_tma_sfb_tensor, work_tile) if cutlass.const_expr(self.use_2cta_instrs): if not self.is_leader_cta: load_shift = dynamic_mainloop.compute_non_leader_cta_load_shift( - valid_tokens_in_tile=work_tile.valid_tokens_in_cta_tile, - mma_tiler_n=self.mma_tiler_mnk[1], + valid_tokens_in_tile=work_tile.valid_tokens_in_cta_tile, mma_tiler_n=self.mma_tiler_mnk[1] ) real_b = cute.domain_offset((load_shift, 0, 0), real_b) - global_b = cute.local_tile( - real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None) - ) + global_b = cute.local_tile(real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None)) global_sfb = cute.local_tile( real_sfb, cute.slice_(self.mma_tiler_sfb, (0, None, None)), (None, None, None) ) @@ -1094,11 +952,7 @@ def run_tma_b( tma_desc_ptr=b_descriptor_ptr, mcast_mask=b_multicast_mask, ) - cute.copy( - b_issue_atom, - global_b_slice[(None, k_tile_idx)], - partitioned_smem_b[(None, stage_index)], - ) + cute.copy(b_issue_atom, global_b_slice[(None, k_tile_idx)], partitioned_smem_b[(None, stage_index)]) sfb_issue_atom = bind_executable_tma_load_fields( fc1_tma_sfb_exec_atom, tma_bar_ptr=tma_barrier, @@ -1106,9 +960,7 @@ def run_tma_b( mcast_mask=sfb_multicast_mask, ) cute.copy( - sfb_issue_atom, - global_sfb_slice[(None, k_tile_idx)], - partitioned_smem_sfb[(None, stage_index)], + sfb_issue_atom, global_sfb_slice[(None, k_tile_idx)], partitioned_smem_sfb[(None, stage_index)] ) else: iket.range_push("tma_token_fc2") @@ -1124,24 +976,17 @@ def run_tma_b( iket.range_pop() k_tile_count = self.fc2_k_tile_count - real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor( - "b", fc2_tma_b_tensor, work_tile - ) - real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfb", fc2_tma_sfb_tensor, work_tile - ) + real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor("b", fc2_tma_b_tensor, work_tile) + real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor("sfb", fc2_tma_sfb_tensor, work_tile) if cutlass.const_expr(self.use_2cta_instrs): if not self.is_leader_cta: load_shift = dynamic_mainloop.compute_non_leader_cta_load_shift( - valid_tokens_in_tile=work_tile.valid_tokens_in_cta_tile, - mma_tiler_n=self.mma_tiler_mnk[1], + valid_tokens_in_tile=work_tile.valid_tokens_in_cta_tile, mma_tiler_n=self.mma_tiler_mnk[1] ) real_b = cute.domain_offset((load_shift, 0, 0), real_b) - global_b = cute.local_tile( - real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None) - ) + global_b = cute.local_tile(real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None)) global_sfb = cute.local_tile( real_sfb, cute.slice_(self.mma_tiler_sfb, (0, None, None)), (None, None, None) ) @@ -1185,11 +1030,7 @@ def run_tma_b( tma_desc_ptr=b_descriptor_ptr, mcast_mask=b_multicast_mask, ) - cute.copy( - b_issue_atom, - global_b_slice[(None, k_tile_idx)], - partitioned_smem_b[(None, stage_index)], - ) + cute.copy(b_issue_atom, global_b_slice[(None, k_tile_idx)], partitioned_smem_b[(None, stage_index)]) sfb_issue_atom = bind_executable_tma_load_fields( fc2_tma_sfb_exec_atom, tma_bar_ptr=tma_barrier, @@ -1197,9 +1038,7 @@ def run_tma_b( mcast_mask=sfb_multicast_mask, ) cute.copy( - sfb_issue_atom, - global_sfb_slice[(None, k_tile_idx)], - partitioned_smem_sfb[(None, stage_index)], + sfb_issue_atom, global_sfb_slice[(None, k_tile_idx)], partitioned_smem_sfb[(None, stage_index)] ) iket.range_pop() @@ -1224,33 +1063,24 @@ def run_mma( tmem_allocator.wait_for_alloc() accumulator_pointer = tmem_allocator.retrieve_ptr(self.acc_dtype) - accumulator_layout = cute.make_layout( - self.accumulator_shape, stride=self.accumulator_stride - ) + accumulator_layout = cute.make_layout(self.accumulator_shape, stride=self.accumulator_stride) accumulator_base = cute.make_tensor(accumulator_pointer, accumulator_layout) sfa_stage_smem_layout = cute.slice_( - blockscaled_utils.make_smem_layout_sfa( - tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1 - ), + blockscaled_utils.make_smem_layout_sfa(tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1), (None, None, None, 0), ) sfb_stage_smem_layout = cute.slice_( - blockscaled_utils.make_smem_layout_sfb( - tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1 - ), + blockscaled_utils.make_smem_layout_sfb(tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1), (None, None, None, 0), ) - sfa_pointer = cute.recast_ptr( - accumulator_pointer + self.num_accumulator_tmem_cols, dtype=self.sf_dtype - ) + sfa_pointer = cute.recast_ptr(accumulator_pointer + self.num_accumulator_tmem_cols, dtype=self.sf_dtype) sfa_layout = blockscaled_utils.make_tmem_layout_sfa( tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, sfa_stage_smem_layout ) tmem_sfa = cute.make_tensor(sfa_pointer, sfa_layout) sfb_pointer = cute.recast_ptr( - accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, - dtype=self.sf_dtype, + accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, dtype=self.sf_dtype ) sfb_layout = blockscaled_utils.make_tmem_layout_sfb( tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, sfb_stage_smem_layout @@ -1283,20 +1113,13 @@ def run_mma( if cutlass.const_expr(self.cta_tile_n == 64): sfb_shift = (work_tile.tile_n_idx % Int32(2)) * Int32(2) shifted_sfb_pointer = cute.recast_ptr( - accumulator_pointer - + self.num_accumulator_tmem_cols - + self.num_sfa_tmem_cols - + sfb_shift, + accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols + sfb_shift, dtype=self.sf_dtype, ) mma_tmem_sfb = cute.make_tensor(shifted_sfb_pointer, sfb_layout) use_n128_sfb_mapping = False - if cutlass.const_expr( - self.quant_kind == QuantKind.nvfp4 and self.cta_tile_n == 256 - ): - aligned_instruction_n = ( - work_tile.valid_tokens_in_cta_tile + Int32(15) - ) & Int32(-16) + if cutlass.const_expr(self.quant_kind == QuantKind.nvfp4 and self.cta_tile_n == 256): + aligned_instruction_n = (work_tile.valid_tokens_in_cta_tile + Int32(15)) & Int32(-16) use_n128_sfb_mapping = aligned_instruction_n <= Int32(128) acc_pipeline.producer_acquire(acc_producer_state) @@ -1311,35 +1134,17 @@ def run_mma( next_full_status = ab_pipeline.consumer_try_wait(next_ab_pipeline_state) s2t_stage_coord = (None, None, None, None, ab_pipeline_state.index) - cute.copy( - tiled_copy_sfa, - partitioned_smem_sfa[s2t_stage_coord], - partitioned_tmem_sfa, - ) - for instruction_index in cutlass.range_constexpr( - self.num_mma_instructions_per_ab_stage - ): + cute.copy(tiled_copy_sfa, partitioned_smem_sfa[s2t_stage_coord], partitioned_tmem_sfa) + for instruction_index in cutlass.range_constexpr(self.num_mma_instructions_per_ab_stage): # Rest_Tiler order: N0K0, N128K0, N0K64, N128K64. - lower_n_lower_k_source_coord = ( - None, - 0, - None, - instruction_index, - ab_pipeline_state.index, - ) + lower_n_lower_k_source_coord = (None, 0, None, instruction_index, ab_pipeline_state.index) lower_n_lower_k_destination_coord = (None, 0, None, instruction_index) cute.copy( tiled_copy_sfb, partitioned_smem_sfb[lower_n_lower_k_source_coord], partitioned_tmem_sfb[lower_n_lower_k_destination_coord], ) - lower_n_upper_k_source_coord = ( - None, - 2, - None, - instruction_index, - ab_pipeline_state.index, - ) + lower_n_upper_k_source_coord = (None, 2, None, instruction_index, ab_pipeline_state.index) compact_upper_k_destination_coord = (None, 1, None, instruction_index) cute.copy( tiled_copy_sfb, @@ -1378,16 +1183,8 @@ def run_mma( next_full_status = ab_pipeline.consumer_try_wait(next_ab_pipeline_state) s2t_stage_coord = (None, None, None, None, ab_pipeline_state.index) - cute.copy( - tiled_copy_sfa, - partitioned_smem_sfa[s2t_stage_coord], - partitioned_tmem_sfa, - ) - cute.copy( - tiled_copy_sfb, - partitioned_smem_sfb[s2t_stage_coord], - partitioned_tmem_sfb, - ) + cute.copy(tiled_copy_sfa, partitioned_smem_sfa[s2t_stage_coord], partitioned_tmem_sfa) + cute.copy(tiled_copy_sfb, partitioned_smem_sfb[s2t_stage_coord], partitioned_tmem_sfb) a_stage_frag = fragment_a[(None, None, None, ab_pipeline_state.index)] b_stage_frag = fragment_b[(None, None, None, ab_pipeline_state.index)] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_mainloop_gen_specialized.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_mainloop_gen_specialized.py index 9533ff4ff895..029cd8a08ee9 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_mainloop_gen_specialized.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_mainloop_gen_specialized.py @@ -1,7 +1,5 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: BSD-3-Clause -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: BSD-3-Clause """Mainloop component for the Rubin generation-phase block-scaled swap-AB FC12 kernel. Three properties separate this from the two mainloops it descends from. @@ -54,31 +52,23 @@ from .....quant_def import QuantKind from ....schedulers.base import SchedulerConsumer from ....schedulers.fc12_mapping import BlockPhase -from ..local_mega.tma_gather import sm107_tma_gather4_load from . import dynamic_mainloop from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension +from .tma_gather import sm107_tma_gather4_load class BlockScaledSwapAbFc12MainloopGenSpecialized(KernelComponent): """Own all Rubin load, MMA, pipeline, SMEM, and TMEM mainloop state.""" - ab_pipeline_mbarriers_region: ClassVar[str] = ( - "rubin.genphase_fc12.mainloop.ab_pipeline_mbarriers" - ) - token_pipeline_mbarriers_region: ClassVar[str] = ( - "rubin.genphase_fc12.mainloop.token_pipeline_mbarriers" - ) + ab_pipeline_mbarriers_region: ClassVar[str] = "rubin.genphase_fc12.mainloop.ab_pipeline_mbarriers" + token_pipeline_mbarriers_region: ClassVar[str] = "rubin.genphase_fc12.mainloop.token_pipeline_mbarriers" a_smem_tensor_region: ClassVar[str] = "rubin.genphase_fc12.mainloop.a_smem_tensor" b_smem_tensor_region: ClassVar[str] = "rubin.genphase_fc12.mainloop.b_smem_tensor" sfa_smem_tensor_region: ClassVar[str] = "rubin.genphase_fc12.mainloop.sfa_smem_tensor" sfb_smem_tensor_region: ClassVar[str] = "rubin.genphase_fc12.mainloop.sfb_smem_tensor" - acc_pipeline_mbarriers_region: ClassVar[str] = ( - "rubin.genphase_fc12.mainloop.acc_pipeline_mbarriers" - ) + acc_pipeline_mbarriers_region: ClassVar[str] = "rubin.genphase_fc12.mainloop.acc_pipeline_mbarriers" tmem_holding_buffer_region: ClassVar[str] = "rubin.genphase_fc12.mainloop.tmem_holding_buffer" - tmem_deallocation_mbarrier_region: ClassVar[str] = ( - "rubin.genphase_fc12.mainloop.tmem_deallocation_mbarrier" - ) + tmem_deallocation_mbarrier_region: ClassVar[str] = "rubin.genphase_fc12.mainloop.tmem_deallocation_mbarrier" tmem_allocation_barrier_id: ClassVar[int] = 2 fc1_gather_sync_barrier_id: ClassVar[int] = 3 # Four warps rather than one because the gather4 row operands live in uniform @@ -113,7 +103,8 @@ class BlockScaledSwapAbFc12MainloopGenSpecialized(KernelComponent): # profiling put it -- the token operand is read out of L2, so a handful of # stages covers its latency, but not as few as the L2 latency alone suggests. asymmetric_token_stages: ClassVar[int] = 6 - # The token tile this specialization applies to; other tiles use the shared plan. + # The token tile this specialization applies to. At 128 the shared plan + # already reaches 93% of DRAM peak and is left alone. asymmetric_token_tile: ClassVar[int] = 256 @classmethod @@ -153,9 +144,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.use_2cta_instrs = impl_desc["use_2cta_instrs"] self.tmem_plan = impl_desc["tmem_plan"] self.mainloop_smem_budget_bytes = impl_desc["mainloop_smem_budget_bytes"] - self.num_accumulator_consumer_warps_per_cta = impl_desc[ - "num_accumulator_consumer_warps_per_cta" - ] + self.num_accumulator_consumer_warps_per_cta = impl_desc["num_accumulator_consumer_warps_per_cta"] self.architecture = "sm_107" self.mma_cta_count = 2 if self.use_2cta_instrs else 1 @@ -174,19 +163,11 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.num_accumulator_stages = tmem_plan.accumulator_stage_count self.num_accumulator_pipeline_stages = tmem_plan.accumulator_pipeline_stages if tmem_plan.accumulator_stage_stride_columns != tmem_plan.accumulator_stage_columns: - raise ValueError( - "Rubin Genphase MegaMoE does not support overlapping accumulator stages." - ) + raise ValueError("Rubin Genphase MegaMoE does not support overlapping accumulator stages.") if tmem_plan.accumulator_pipeline_stages != tmem_plan.accumulator_stage_count: - raise ValueError( - "Rubin Genphase MegaMoE requires one pipeline stage per disjoint accumulator stage." - ) + raise ValueError("Rubin Genphase MegaMoE requires one pipeline stage per disjoint accumulator stage.") self.num_accumulator_tmem_cols = tmem_plan.accumulator_columns - self.accumulator_shape = ( - self.cta_tile_m, - self.cta_tile_n, - tmem_plan.accumulator_stage_count, - ) + self.accumulator_shape = (self.cta_tile_m, self.cta_tile_n, tmem_plan.accumulator_stage_count) self.accumulator_stride = (1 << 16, 1, tmem_plan.accumulator_stage_stride_columns) self.mma_instruction = Tcgen05MmaInstruction( @@ -199,12 +180,8 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: sfb_type=self.sf_dtype, sf_vec_size=self.sf_vec_size, ) - self.a_smem_alloc_dtype = tcgen05_smem_alloc_type( - self.a_dtype, self.b_dtype, self.architecture - ) - self.b_smem_alloc_dtype = tcgen05_smem_alloc_type( - self.b_dtype, self.a_dtype, self.architecture - ) + self.a_smem_alloc_dtype = tcgen05_smem_alloc_type(self.a_dtype, self.b_dtype, self.architecture) + self.b_smem_alloc_dtype = tcgen05_smem_alloc_type(self.b_dtype, self.a_dtype, self.architecture) self.sfb_instruction_shape_mnk = ( self.mma_instruction_mnk[0] // self.mma_cta_count, round_up(self.mma_instruction_mnk[1], 128), @@ -220,11 +197,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: sfb_type=self.sf_dtype, sf_vec_size=self.sf_vec_size, ) - self.mma_tiler_sfb = ( - self.sfb_instruction_shape_mnk[0], - self.sfb_instruction_shape_mnk[1], - self.mma_tile_k, - ) + self.mma_tiler_sfb = (self.sfb_instruction_shape_mnk[0], self.sfb_instruction_shape_mnk[1], self.mma_tile_k) self.cluster_layout_shape_vmnk = ( (self.mma_cta_count,), self.cluster_shape_mn[0] // self.mma_cta_count, @@ -237,19 +210,15 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.is_a_mcast = self.num_mcast_ctas_a > 1 self.is_b_mcast = self.num_mcast_ctas_b > 1 - # Only the selected token tile uses asymmetric operand stages. + # Only the 256-wide token tile is specialized: at 128 the shared plan + # already reaches 93% of DRAM peak, so it stays byte-for-byte as it was. self.uses_asymmetric_ab_stages = self.cta_tile_n == self.asymmetric_token_tile self._select_ab_stage_plan() self.num_tma_load_bytes = self.ab_stage_tma_bytes * self.mma_cta_count if self.num_tma_load_bytes != ( - self.a_stage_tma_bytes - + self.b_stage_tma_bytes - + self.sfa_stage_tma_bytes - + self.sfb_stage_tma_bytes + self.a_stage_tma_bytes + self.b_stage_tma_bytes + self.sfa_stage_tma_bytes + self.sfb_stage_tma_bytes ): - raise ValueError( - "Genphase MegaMoE AB stage transaction byte accounting is inconsistent." - ) + raise ValueError("Genphase MegaMoE AB stage transaction byte accounting is inconsistent.") # Under the split plan each pipeline has a single producer whose expectation # is exactly the bytes it issues. The shared plan instead has two producers # splitting one combined expectation, so neither share matches the bytes @@ -258,14 +227,11 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.token_pipeline_tx_count = self.b_stage_tma_bytes + self.sfb_stage_tma_bytes if not self.uses_asymmetric_ab_stages: if self.num_tma_load_bytes % 2 != 0: - raise ValueError( - "Genphase MegaMoE total AB transaction bytes must split evenly across two arrivals." - ) + raise ValueError("Genphase MegaMoE total AB transaction bytes must split evenly across two arrivals.") self.ab_pipeline_tx_count_per_producer = self.num_tma_load_bytes // 2 self._plan_gather_geometry() self.fc1_gather_sync_barrier = pipeline.NamedBarrier( - barrier_id=self.fc1_gather_sync_barrier_id, - num_threads=32 * (1 + self.num_gather_b_warps), + barrier_id=self.fc1_gather_sync_barrier_id, num_threads=32 * (1 + self.num_gather_b_warps) ) def _plan_gather_geometry(self) -> None: @@ -291,9 +257,7 @@ def _plan_gather_geometry(self) -> None: if self.gather_b_rows_per_cta % run_span_rows != 0: # The MMA CTA count cancels out of both sides, so the requirement is # purely on tile N and cluster M. - required_multiple = ( - self.gather_index_load_rows * self.num_gather_b_warps * self.cluster_shape_mn[0] - ) + required_multiple = self.gather_index_load_rows * self.num_gather_b_warps * self.cluster_shape_mn[0] raise ValueError( f"Genphase MegaMoE gather runs must divide evenly across {self.gather_issuing_units} issuing " f"units: tile N {self.cta_tile_n} must be a multiple of {required_multiple}." @@ -335,10 +299,7 @@ def _validate_configuration(self) -> None: ): if len(dimensions) != expected_rank: raise ValueError(f"{field_name} must contain {expected_rank} dimensions.") - if not all( - isinstance(dimension, int) and not isinstance(dimension, bool) - for dimension in dimensions - ): + if not all(isinstance(dimension, int) and not isinstance(dimension, bool) for dimension in dimensions): raise TypeError(f"{field_name} dimensions must be Python integers.") if any(dimension <= 0 for dimension in dimensions): raise ValueError(f"{field_name} dimensions must be positive.") @@ -355,9 +316,7 @@ def _validate_configuration(self) -> None: f"got {self.mma_instruction_mnk[0]}." ) if self.mma_instruction_mnk[:2] != self.mma_tiler_mnk[:2]: - raise NotImplementedError( - "Rubin Genphase MegaMoE does not implement M/N instruction repetition." - ) + raise NotImplementedError("Rubin Genphase MegaMoE does not implement M/N instruction repetition.") # Tile N 64 is excluded: it forces the halved SFB tile index and the SFB # TMEM shift, two special paths whose interaction with a dynamic # instruction N is not worth carrying for a token extent this small. @@ -382,13 +341,9 @@ def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: # Under the shared plan this one array serves all four operands. Under the # split plan it becomes the weight side's and the token side gets its own, # because the two depths no longer index the same stage. - smem_workspace.register_mbarrier( - self.ab_pipeline_mbarriers_region, self.num_weight_ab_stages * 2 - ) + smem_workspace.register_mbarrier(self.ab_pipeline_mbarriers_region, self.num_weight_ab_stages * 2) if self.uses_asymmetric_ab_stages: - smem_workspace.register_mbarrier( - self.token_pipeline_mbarriers_region, self.num_token_ab_stages * 2 - ) + smem_workspace.register_mbarrier(self.token_pipeline_mbarriers_region, self.num_token_ab_stages * 2) for region_name, region in ( (self.a_smem_tensor_region, self.a_smem_region), (self.b_smem_tensor_region, self.b_smem_region), @@ -403,37 +358,23 @@ def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: swizzle=region.swizzle, byte_alignment=region.byte_alignment, ) - smem_workspace.register_mbarrier( - self.acc_pipeline_mbarriers_region, self.num_accumulator_pipeline_stages * 2 - ) + smem_workspace.register_mbarrier(self.acc_pipeline_mbarriers_region, self.num_accumulator_pipeline_stages * 2) smem_workspace.register_mbarrier(self.tmem_deallocation_mbarrier_region, 1) - smem_workspace.register_tensor( - self.tmem_holding_buffer_region, cutlass.Int32, (1,), byte_alignment=8 - ) + smem_workspace.register_tensor(self.tmem_holding_buffer_region, cutlass.Int32, (1,), byte_alignment=8) def _select_ab_stage_plan(self) -> None: stage_one_regions = make_smem_layouts( - self.mma_instruction, - self.mma_tiler_mnk, - 1, - (self.a_major_mode, self.b_major_mode), - self.architecture, + self.mma_instruction, self.mma_tiler_mnk, 1, (self.a_major_mode, self.b_major_mode), self.architecture ) independent_sfb_region = make_smem_layouts( - self.sfb_mma_instruction, - self.mma_tiler_sfb, - 1, - (self.a_major_mode, self.b_major_mode), - self.architecture, + self.sfb_mma_instruction, self.mma_tiler_sfb, 1, (self.a_major_mode, self.b_major_mode), self.architecture )[3] if ( independent_sfb_region.shape != stage_one_regions[3].shape or independent_sfb_region.stride != stage_one_regions[3].stride or independent_sfb_region.nbytes != stage_one_regions[3].nbytes ): - raise ValueError( - "The Rubin independent SFB plan does not match the primary MMA SFB operand." - ) + raise ValueError("The Rubin independent SFB plan does not match the primary MMA SFB operand.") def packed_bytes(region, dtype) -> int: return (region.cosize * int(dtype.width) + 7) // 8 @@ -448,24 +389,17 @@ def packed_bytes(region, dtype) -> int: stage_one_regions[3].nbytes, ) self.ab_stage_tma_bytes = sum(stage_operand_tma_bytes) - ( - self.a_stage_tma_bytes, - self.b_stage_tma_bytes, - self.sfa_stage_tma_bytes, - self.sfb_stage_tma_bytes, - ) = tuple(stage_bytes * self.mma_cta_count for stage_bytes in stage_operand_tma_bytes) + (self.a_stage_tma_bytes, self.b_stage_tma_bytes, self.sfa_stage_tma_bytes, self.sfb_stage_tma_bytes) = tuple( + stage_bytes * self.mma_cta_count for stage_bytes in stage_operand_tma_bytes + ) mbarrier_bytes = int(cutlass.Int64.width) // 8 plan_tail_bytes = (2 * self.num_accumulator_pipeline_stages + 2) * mbarrier_bytes stage_budget_bytes = self.mainloop_smem_budget_bytes - plan_tail_bytes # Each pipeline spends two mbarriers per stage, full and empty. The shared # plan pays that once for all four operands; the split plan pays it twice, # which is the only overhead the split introduces. - weight_stage_cost_bytes = ( - stage_one_regions[0].nbytes + stage_one_regions[2].nbytes + 2 * mbarrier_bytes - ) - token_stage_cost_bytes = ( - stage_one_regions[1].nbytes + stage_one_regions[3].nbytes + 2 * mbarrier_bytes - ) + weight_stage_cost_bytes = stage_one_regions[0].nbytes + stage_one_regions[2].nbytes + 2 * mbarrier_bytes + token_stage_cost_bytes = stage_one_regions[1].nbytes + stage_one_regions[3].nbytes + 2 * mbarrier_bytes if self.uses_asymmetric_ab_stages: self.num_token_ab_stages = self.asymmetric_token_stages @@ -516,12 +450,7 @@ def packed_bytes(region, dtype) -> int: ) self.selected_smem_bytes = sum( region.nbytes - for region in ( - self.a_smem_region, - self.b_smem_region, - self.sfa_smem_region, - self.sfb_smem_region, - ) + for region in (self.a_smem_region, self.b_smem_region, self.sfa_smem_region, self.sfb_smem_region) ) def make_tiled_mma(self) -> cute.TiledMma: @@ -591,11 +520,8 @@ def materialize_codegen_members(self) -> None: ("SFA", self.sfa_smem_layout_staged, self.sfa_smem_region), ("SFB", self.sfb_smem_layout_staged, self.sfb_smem_region), ): - if ( - materialized_layout.shape != planned_region.shape - or not strides_equal_ignoring_singletons( - planned_region.shape, materialized_layout.stride, planned_region.stride - ) + if materialized_layout.shape != planned_region.shape or not strides_equal_ignoring_singletons( + planned_region.shape, materialized_layout.stride, planned_region.stride ): raise ValueError( f"{tensor_name} SMEM plan mismatch: Python " @@ -621,24 +547,12 @@ def prepare_tma_load_params( sfa_stage_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0)) sfb_stage_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0)) - a_operation = sm100_utils.cluster_shape_to_tma_atom_A( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - b_operation = sm100_utils.cluster_shape_to_tma_atom_B( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - sfa_operation = sm100_utils.cluster_shape_to_tma_atom_A( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - sfb_operation = sm100_utils.cluster_shape_to_tma_atom_SFB( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - a_internal_type = ( - self.a_smem_alloc_dtype if self.a_smem_alloc_dtype is not self.a_dtype else None - ) - b_internal_type = ( - self.b_smem_alloc_dtype if self.b_smem_alloc_dtype is not self.b_dtype else None - ) + a_operation = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, self.tiled_mma.thr_id) + b_operation = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, self.tiled_mma.thr_id) + sfa_operation = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, self.tiled_mma.thr_id) + sfb_operation = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, self.tiled_mma.thr_id) + a_internal_type = self.a_smem_alloc_dtype if self.a_smem_alloc_dtype is not self.a_dtype else None + b_internal_type = self.b_smem_alloc_dtype if self.b_smem_alloc_dtype is not self.b_dtype else None fc1_a_atom, fc1_a_tensor = cute.nvgpu.make_tiled_tma_atom_A( a_operation, @@ -656,9 +570,7 @@ def prepare_tma_load_params( # gather supplies coordinates directly and never partitions a tensor view. fc1_b_2d = fc1_b[(None, None, 0)] gather_b_base = cute.make_layout((1, self.mma_tile_k), stride=(self.mma_tile_k, 1)) - gather_b_smem_layout = cute.make_composed_layout( - cute.make_swizzle(3, 4, 3), 0, gather_b_base - ) + gather_b_smem_layout = cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, gather_b_base) fc1_b_atom, _ = cpasync.make_tiled_tma_atom( cpasync.CopyBulkTensorTileG2SOp(), fc1_b_2d, gather_b_smem_layout, (1, self.mma_tile_k) ) @@ -765,17 +677,10 @@ def assign_device_members( ) self.fc1_k_tile_count = ceil_div(bound_hidden_size, self.mma_tile_k) self.fc2_k_tile_count = ceil_div(bound_intermediate_gateup_size // 2, self.mma_tile_k) - self.sfb_vmnk_coord = ( - 0, - cta_coord_in_cluster[0], - cta_coord_in_cluster[1], - cta_coord_in_cluster[2], - ) + self.sfb_vmnk_coord = (0, cta_coord_in_cluster[0], cta_coord_in_cluster[1], cta_coord_in_cluster[2]) @cute.jit - def create_ab_pipeline( - self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer - ) -> pipeline.PipelineAsync: + def create_ab_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> pipeline.PipelineAsync: """Create the shared two-producer TMA-to-UMMA pipeline. Each producer arrival contributes the same accounting share. The two @@ -798,9 +703,7 @@ def create_ab_pipeline( ) @cute.jit - def create_weight_pipeline( - self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer - ) -> pipeline.PipelineAsync: + def create_weight_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> pipeline.PipelineAsync: """Create the weight half of the split TMA-to-UMMA pipeline. One producer, TMA-A, and its expectation is exactly the A and SFA bytes it @@ -821,9 +724,7 @@ def create_weight_pipeline( ) @cute.jit - def create_token_pipeline( - self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer - ) -> pipeline.PipelineAsync: + def create_token_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> pipeline.PipelineAsync: """Create the token half of the split TMA-to-UMMA pipeline. One producer, TMA-B, registering the B and SFB bytes. Gather-B still never @@ -848,9 +749,7 @@ def create_token_pipeline( def create_acc_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer): """Create the UMMA-to-epilogue accumulator pipeline.""" producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) - consumer_thread_count = ( - self.num_accumulator_consumer_warps_per_cta * 32 * self.mma_cta_count - ) + consumer_thread_count = self.num_accumulator_consumer_warps_per_cta * 32 * self.mma_cta_count consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, consumer_thread_count) return pipeline.PipelineUmmaAsync.create( barrier_storage=smem_workspace.ptr(self.acc_pipeline_mbarriers_region, smem_base), @@ -862,9 +761,7 @@ def create_acc_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Poi ) @cute.jit - def create_tmem_allocator( - self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer, *, allocator_warp_id: int - ): + def create_tmem_allocator(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer, *, allocator_warp_id: int): """Bind the SM107 TMEM allocator to finalized SMEM regions.""" allocation_barrier = pipeline.NamedBarrier( barrier_id=self.tmem_allocation_barrier_id, @@ -875,9 +772,7 @@ def create_tmem_allocator( barrier_for_retrieve=allocation_barrier, allocator_warp_id=allocator_warp_id, is_two_cta=self.use_2cta_instrs, - two_cta_tmem_dealloc_mbar_ptr=smem_workspace.ptr( - self.tmem_deallocation_mbarrier_region, smem_base - ), + two_cta_tmem_dealloc_mbar_ptr=smem_workspace.ptr(self.tmem_deallocation_mbarrier_region, smem_base), arch=self.architecture, ) @@ -891,18 +786,12 @@ def _s2t_copy_and_partition(self, smem_tensor: cute.Tensor, tmem_tensor: cute.Te mn_mode = cute.get(compact_smem_tensor.layout, mode=[0, 0]) mn_mode = cute.append(mn_mode, cute.make_layout((4,), stride=(0,))) - broadcast_layout = cute.append( - cute.group_modes(mn_mode, 0), cute.get(compact_smem_tensor.layout, mode=[0, 1]) - ) + broadcast_layout = cute.append(cute.group_modes(mn_mode, 0), cute.get(compact_smem_tensor.layout, mode=[0, 1])) broadcast_layout = cute.append( cute.group_modes(broadcast_layout, 0), cute.get(compact_smem_tensor.layout, mode=[1]) ) - broadcast_layout = cute.append( - broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[2]) - ) - broadcast_layout = cute.append( - broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[3]) - ) + broadcast_layout = cute.append(broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[2])) + broadcast_layout = cute.append(broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[3])) broadcast_smem_tensor = cute.make_tensor(compact_smem_tensor.iterator, broadcast_layout) partitioned_smem = thread_copy.partition_S(broadcast_smem_tensor) @@ -931,17 +820,10 @@ def run_tma_a( multicast_mask = None if cutlass.const_expr(self.is_a_mcast or self.use_2cta_instrs): multicast_mask = tma_multicast_mask( - self.cluster_shape_mn, - None, - self.cta_coord_in_cluster, - None, - self.use_2cta_instrs, - "a", + self.cluster_shape_mn, None, self.cta_coord_in_cluster, None, self.use_2cta_instrs, "a" ) - a_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_vmnk, (0, 0, None, 0)).shape - ) + a_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_vmnk, (0, 0, None, 0)).shape) sfa_cta_layout = a_cta_layout thread_mma = self.tiled_mma.get_slice(self.mma_cta_index) @@ -954,16 +836,10 @@ def run_tma_a( if is_fc1: iket.range_push("tma_weight_fc1") k_tile_count = self.fc1_k_tile_count - real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor( - "a", fc1_tma_a_tensor, work_tile - ) - real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfa", fc1_tma_sfa_tensor, work_tile - ) + real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor("a", fc1_tma_a_tensor, work_tile) + real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor("sfa", fc1_tma_sfa_tensor, work_tile) - global_a = cute.local_tile( - real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) - ) + global_a = cute.local_tile(real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None)) global_sfa = cute.local_tile( real_sfa, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) ) @@ -1018,16 +894,10 @@ def run_tma_a( else: iket.range_push("tma_weight_fc2") k_tile_count = self.fc2_k_tile_count - real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor( - "a", fc2_tma_a_tensor, work_tile - ) - real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfa", fc2_tma_sfa_tensor, work_tile - ) + real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor("a", fc2_tma_a_tensor, work_tile) + real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor("sfa", fc2_tma_sfa_tensor, work_tile) - global_a = cute.local_tile( - real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) - ) + global_a = cute.local_tile(real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None)) global_sfa = cute.local_tile( real_sfa, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) ) @@ -1120,28 +990,14 @@ def run_tma_b( sfb_multicast_mask = None if cutlass.const_expr(self.is_b_mcast or self.use_2cta_instrs): b_multicast_mask = tma_multicast_mask( - self.cluster_shape_mn, - None, - self.cta_coord_in_cluster, - None, - self.use_2cta_instrs, - "b", + self.cluster_shape_mn, None, self.cta_coord_in_cluster, None, self.use_2cta_instrs, "b" ) sfb_multicast_mask = tma_multicast_mask( - self.cluster_shape_mn, - None, - self.cta_coord_in_cluster, - None, - self.use_2cta_instrs, - "sfb", + self.cluster_shape_mn, None, self.cta_coord_in_cluster, None, self.use_2cta_instrs, "sfb" ) - b_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_vmnk, (0, None, 0, 0)).shape - ) - sfb_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape - ) + b_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_vmnk, (0, None, 0, 0)).shape) + sfb_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape) thread_mma = self.tiled_mma.get_slice(self.mma_cta_index) thread_mma_sfb = self.tiled_mma_sfb.get_slice(self.mma_cta_index) @@ -1153,16 +1009,10 @@ def run_tma_b( iket.range_push("tma_b.wait_expert_sf_ready") expert_idx = work_tile.expert_idx expert_sf_blocks = ceil_div(expert_sizes[expert_idx], Int32(self.sf_padding_block)) - spin_wait( - fc1_sf_ready_pointer + expert_idx, - lambda value: value == expert_sf_blocks, - scope="gpu", - ) + spin_wait(fc1_sf_ready_pointer + expert_idx, lambda value: value == expert_sf_blocks, scope="gpu") iket.range_pop() k_tile_count = self.fc1_k_tile_count - real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfb", fc1_tma_sfb_tensor, work_tile - ) + real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor("sfb", fc1_tma_sfb_tensor, work_tile) global_sfb = cute.local_tile( real_sfb, cute.slice_(self.mma_tiler_sfb, (0, None, None)), (None, None, None) @@ -1216,12 +1066,8 @@ def run_tma_b( iket.range_pop() k_tile_count = self.fc2_k_tile_count - real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor( - "b", fc2_tma_b_tensor, work_tile - ) - real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfb", fc2_tma_sfb_tensor, work_tile - ) + real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor("b", fc2_tma_b_tensor, work_tile) + real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor("sfb", fc2_tma_sfb_tensor, work_tile) # Under two-CTA instructions the dynamic N window is split across # the pair, so the non-leader's rows start at the runtime midpoint @@ -1229,14 +1075,11 @@ def run_tma_b( if cutlass.const_expr(self.use_2cta_instrs): if not self.is_leader_cta: load_shift = dynamic_mainloop.compute_non_leader_cta_load_shift( - valid_tokens_in_tile=work_tile.valid_tokens_in_cta_tile, - mma_tiler_n=self.mma_tiler_mnk[1], + valid_tokens_in_tile=work_tile.valid_tokens_in_cta_tile, mma_tiler_n=self.mma_tiler_mnk[1] ) real_b = cute.domain_offset((load_shift, 0, 0), real_b) - global_b = cute.local_tile( - real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None) - ) + global_b = cute.local_tile(real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None)) global_sfb = cute.local_tile( real_sfb, cute.slice_(self.mma_tiler_sfb, (0, None, None)), (None, None, None) ) @@ -1315,9 +1158,7 @@ def run_tma_gather_b1( producer. FC2 is fully produced by TMA-A and TMA-B, so on those tiles this warp only keeps its register-held pipeline state aligned. """ - row_ids = cute.make_rmem_tensor( - cute.make_layout((self.gather_b_rows_per_unit,)), cutlass.Int32 - ) + row_ids = cute.make_rmem_tensor(cute.make_layout((self.gather_b_rows_per_unit,)), cutlass.Int32) row_id_runs = cute.zipped_divide(row_ids, (self.gather_index_load_rows,)) index_load_atom = cute.make_copy_atom( cute.nvgpu.CopyUniversalOp(), @@ -1328,12 +1169,7 @@ def run_tma_gather_b1( gather_b_multicast_mask = None if cutlass.const_expr(self.is_b_mcast or self.use_2cta_instrs): gather_b_multicast_mask = tma_multicast_mask( - self.cluster_shape_mn, - None, - self.cta_coord_in_cluster, - None, - self.use_2cta_instrs, - "b", + self.cluster_shape_mn, None, self.cta_coord_in_cluster, None, self.use_2cta_instrs, "b" ) issuing_unit = self.main_vmnk_coord[1] * Int32(self.num_gather_b_warps) + gather_warp_idx @@ -1342,10 +1178,7 @@ def run_tma_gather_b1( is_fc1 = work_tile.phase == Int32(BlockPhase.Linear1) if is_fc1: iket.range_push("tma_gather_b1_fc1") - tile_pool_base = ( - work_tile.cumulative_data_physical_row - + work_tile.tile_n_idx * Int32(self.cta_tile_n) - ) + tile_pool_base = work_tile.cumulative_data_physical_row + work_tile.tile_n_idx * Int32(self.cta_tile_n) # The MMA issues a dynamic instruction N, so under two-CTA # instructions the pair splits that runtime extent rather than the # static tile: the non-leader's local row zero is the midpoint of @@ -1353,12 +1186,9 @@ def run_tma_gather_b1( # `compute_non_leader_cta_load_shift` does for the tiled TMA legs. # Rounding to sixteen keeps that midpoint a multiple of eight, so # every run stays aligned to its 256-bit index load. - cta_row_base = self.mma_cta_index * ( - round_up(work_tile.valid_tokens_in_cta_tile, 16) // Int32(2) - ) + cta_row_base = self.mma_cta_index * (round_up(work_tile.valid_tokens_in_cta_tile, 16) // Int32(2)) index_runs = cute.zipped_divide( - cute.domain_offset((tile_pool_base + cta_row_base,), gather_index), - (self.gather_index_load_rows,), + cute.domain_offset((tile_pool_base + cta_row_base,), gather_index), (self.gather_index_load_rows,) ) for local_run in cutlass.range_constexpr(self.gather_b_runs_per_unit): run_index = Int32(local_run * self.gather_issuing_units) + issuing_unit @@ -1376,9 +1206,7 @@ def run_tma_gather_b1( tma_barrier = ab_pipeline.producer_get_barrier(ab_pipeline_state) stage_index = ab_pipeline_state.index col_k = k_tile_idx * Int32(self.mma_tile_k) - stage_base_elements = stage_index * Int32( - self.gather_b_rows_per_cta * self.mma_tile_k - ) + stage_base_elements = stage_index * Int32(self.gather_b_rows_per_cta * self.mma_tile_k) for local_run in cutlass.range_constexpr(self.gather_b_runs_per_unit): run_index = Int32(local_run * self.gather_issuing_units) + issuing_unit @@ -1389,10 +1217,7 @@ def run_tma_gather_b1( destination_offset = run_base_elements + Int32( group_in_run * self.gather4_width * self.mma_tile_k ) - row_start = ( - local_run * self.gather_index_load_rows - + group_in_run * self.gather4_width - ) + row_start = local_run * self.gather_index_load_rows + group_in_run * self.gather4_width with cute.arch.elect_one(): sm107_tma_gather4_load( fc1_tma_b_atom, @@ -1437,33 +1262,24 @@ def run_mma( tmem_allocator.wait_for_alloc() accumulator_pointer = tmem_allocator.retrieve_ptr(self.acc_dtype) - accumulator_layout = cute.make_layout( - self.accumulator_shape, stride=self.accumulator_stride - ) + accumulator_layout = cute.make_layout(self.accumulator_shape, stride=self.accumulator_stride) accumulator_base = cute.make_tensor(accumulator_pointer, accumulator_layout) sfa_stage_smem_layout = cute.slice_( - blockscaled_utils.make_smem_layout_sfa( - tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1 - ), + blockscaled_utils.make_smem_layout_sfa(tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1), (None, None, None, 0), ) sfb_stage_smem_layout = cute.slice_( - blockscaled_utils.make_smem_layout_sfb( - tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1 - ), + blockscaled_utils.make_smem_layout_sfb(tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1), (None, None, None, 0), ) - sfa_pointer = cute.recast_ptr( - accumulator_pointer + self.num_accumulator_tmem_cols, dtype=self.sf_dtype - ) + sfa_pointer = cute.recast_ptr(accumulator_pointer + self.num_accumulator_tmem_cols, dtype=self.sf_dtype) sfa_layout = blockscaled_utils.make_tmem_layout_sfa( tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, sfa_stage_smem_layout ) tmem_sfa = cute.make_tensor(sfa_pointer, sfa_layout) sfb_pointer = cute.recast_ptr( - accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, - dtype=self.sf_dtype, + accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, dtype=self.sf_dtype ) sfb_layout = blockscaled_utils.make_tmem_layout_sfb( tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, sfb_stage_smem_layout @@ -1496,12 +1312,8 @@ def run_mma( # window compact, so its two K halves come from a different pair of # SMEM rest coordinates than the full-width case. use_n128_sfb_mapping = False - if cutlass.const_expr( - self.quant_kind == QuantKind.nvfp4 and self.cta_tile_n == 256 - ): - use_n128_sfb_mapping = round_up( - work_tile.valid_tokens_in_cta_tile, 16 - ) <= Int32(128) + if cutlass.const_expr(self.quant_kind == QuantKind.nvfp4 and self.cta_tile_n == 256): + use_n128_sfb_mapping = round_up(work_tile.valid_tokens_in_cta_tile, 16) <= Int32(128) acc_pipeline.producer_acquire(acc_producer_state) if use_n128_sfb_mapping: @@ -1515,35 +1327,17 @@ def run_mma( next_full_status = ab_pipeline.consumer_try_wait(next_ab_pipeline_state) s2t_stage_coord = (None, None, None, None, ab_pipeline_state.index) - cute.copy( - tiled_copy_sfa, - partitioned_smem_sfa[s2t_stage_coord], - partitioned_tmem_sfa, - ) - for instruction_index in cutlass.range_constexpr( - self.num_mma_instructions_per_ab_stage - ): + cute.copy(tiled_copy_sfa, partitioned_smem_sfa[s2t_stage_coord], partitioned_tmem_sfa) + for instruction_index in cutlass.range_constexpr(self.num_mma_instructions_per_ab_stage): # Rest_Tiler order: N0K0, N128K0, N0K64, N128K64. - lower_n_lower_k_source_coord = ( - None, - 0, - None, - instruction_index, - ab_pipeline_state.index, - ) + lower_n_lower_k_source_coord = (None, 0, None, instruction_index, ab_pipeline_state.index) lower_n_lower_k_destination_coord = (None, 0, None, instruction_index) cute.copy( tiled_copy_sfb, partitioned_smem_sfb[lower_n_lower_k_source_coord], partitioned_tmem_sfb[lower_n_lower_k_destination_coord], ) - lower_n_upper_k_source_coord = ( - None, - 2, - None, - instruction_index, - ab_pipeline_state.index, - ) + lower_n_upper_k_source_coord = (None, 2, None, instruction_index, ab_pipeline_state.index) compact_upper_k_destination_coord = (None, 1, None, instruction_index) cute.copy( tiled_copy_sfb, @@ -1582,16 +1376,8 @@ def run_mma( next_full_status = ab_pipeline.consumer_try_wait(next_ab_pipeline_state) s2t_stage_coord = (None, None, None, None, ab_pipeline_state.index) - cute.copy( - tiled_copy_sfa, - partitioned_smem_sfa[s2t_stage_coord], - partitioned_tmem_sfa, - ) - cute.copy( - tiled_copy_sfb, - partitioned_smem_sfb[s2t_stage_coord], - partitioned_tmem_sfb, - ) + cute.copy(tiled_copy_sfa, partitioned_smem_sfa[s2t_stage_coord], partitioned_tmem_sfa) + cute.copy(tiled_copy_sfb, partitioned_smem_sfb[s2t_stage_coord], partitioned_tmem_sfb) a_stage_frag = fragment_a[(None, None, None, ab_pipeline_state.index)] b_stage_frag = fragment_b[(None, None, None, ab_pipeline_state.index)] @@ -1651,33 +1437,24 @@ def run_mma_asymmetric( tmem_allocator.wait_for_alloc() accumulator_pointer = tmem_allocator.retrieve_ptr(self.acc_dtype) - accumulator_layout = cute.make_layout( - self.accumulator_shape, stride=self.accumulator_stride - ) + accumulator_layout = cute.make_layout(self.accumulator_shape, stride=self.accumulator_stride) accumulator_base = cute.make_tensor(accumulator_pointer, accumulator_layout) sfa_stage_smem_layout = cute.slice_( - blockscaled_utils.make_smem_layout_sfa( - tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1 - ), + blockscaled_utils.make_smem_layout_sfa(tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1), (None, None, None, 0), ) sfb_stage_smem_layout = cute.slice_( - blockscaled_utils.make_smem_layout_sfb( - tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1 - ), + blockscaled_utils.make_smem_layout_sfb(tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1), (None, None, None, 0), ) - sfa_pointer = cute.recast_ptr( - accumulator_pointer + self.num_accumulator_tmem_cols, dtype=self.sf_dtype - ) + sfa_pointer = cute.recast_ptr(accumulator_pointer + self.num_accumulator_tmem_cols, dtype=self.sf_dtype) sfa_layout = blockscaled_utils.make_tmem_layout_sfa( tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, sfa_stage_smem_layout ) tmem_sfa = cute.make_tensor(sfa_pointer, sfa_layout) sfb_pointer = cute.recast_ptr( - accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, - dtype=self.sf_dtype, + accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, dtype=self.sf_dtype ) sfb_layout = blockscaled_utils.make_tmem_layout_sfb( tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, sfb_stage_smem_layout @@ -1713,12 +1490,8 @@ def run_mma_asymmetric( # state machine out of the branch: the predicate is warp-uniform # and this loop has no state machine to duplicate. use_n128_sfb_mapping = False - if cutlass.const_expr( - self.quant_kind == QuantKind.nvfp4 and self.cta_tile_n == 256 - ): - use_n128_sfb_mapping = round_up( - work_tile.valid_tokens_in_cta_tile, 16 - ) <= Int32(128) + if cutlass.const_expr(self.quant_kind == QuantKind.nvfp4 and self.cta_tile_n == 256): + use_n128_sfb_mapping = round_up(work_tile.valid_tokens_in_cta_tile, 16) <= Int32(128) acc_pipeline.producer_acquire(acc_producer_state) for k_tile_idx in cutlass.range(k_tile_count, unroll=1): @@ -1726,15 +1499,9 @@ def run_mma_asymmetric( token_pipeline.consumer_wait(token_pipeline_state) weight_stage_coord = (None, None, None, None, weight_pipeline_state.index) - cute.copy( - tiled_copy_sfa, - partitioned_smem_sfa[weight_stage_coord], - partitioned_tmem_sfa, - ) + cute.copy(tiled_copy_sfa, partitioned_smem_sfa[weight_stage_coord], partitioned_tmem_sfa) if use_n128_sfb_mapping: - for instruction_index in cutlass.range_constexpr( - self.num_mma_instructions_per_ab_stage - ): + for instruction_index in cutlass.range_constexpr(self.num_mma_instructions_per_ab_stage): # Rest_Tiler order: N0K0, N128K0, N0K64, N128K64. lower_n_lower_k_source_coord = ( None, @@ -1762,11 +1529,7 @@ def run_mma_asymmetric( ) else: token_stage_coord = (None, None, None, None, token_pipeline_state.index) - cute.copy( - tiled_copy_sfb, - partitioned_smem_sfb[token_stage_coord], - partitioned_tmem_sfb, - ) + cute.copy(tiled_copy_sfb, partitioned_smem_sfb[token_stage_coord], partitioned_tmem_sfb) a_stage_frag = fragment_a[(None, None, None, weight_pipeline_state.index)] b_stage_frag = fragment_b[(None, None, None, token_pipeline_state.index)] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py index b894e963af64..3c5bb77f999c 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py @@ -1,11 +1,11 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Fused pull-dispatch MegaMoE kernel composition for Rubin.""" +import hashlib + from typing import ClassVar, Optional, Tuple import cuda.bindings.driver as cuda + import cutlass import cutlass.cute as cute import cutlass.pipeline as pipeline @@ -14,30 +14,23 @@ from cutlass.cute.nvgpu import OperandMajorMode from cutlass.cutlass_dsl import Boolean, Int32 -from .....api import ( - ImplDesc, - KernelClass, - OptionalRequirement, - ProblemDesc, - StaticOrRuntimeIntegerType, -) +from .....api import ImplDesc, KernelClass, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType from .....communication.nvlink_domain.token_comm import TokenCommArgs, TokenCommNonDeterministic -from .....helpers.cute_py_helpers import ( - Tcgen05MmaInstruction, - make_tcgen05_tmem_plan, - tcgen05_block_scaled_acc_dtype, +from .....communication.nvlink_domain.token_comm_cta_mainloop import TokenCommCtaMainloop +from .....helpers.helper_weight_ready_gate import ( + derive_helper_weight_ready_topology, + make_helper_weight_ready_generation_gate, ) +from .....helpers.cute_py_helpers import Tcgen05MmaInstruction, make_tcgen05_tmem_plan, tcgen05_block_scaled_acc_dtype from .....helpers.device_workspace import DeviceWorkspace +from .....helpers.dlb_expert_count import resolve_dlb_expert_counts +from .....helpers.fc2_ready_geometry import derive_fc2_ready_geometry from .....helpers.iket_compat import iket +from .....helpers.megamoe_aot import make_megamoe_aot_callable from .....helpers.smem_workspace import SmemWorkspace +from .....helpers.topk_reduce import TopkReduce from .....helpers.utils import ceil_div, round_up from .....quant_def import CombineFormat, QuantKind -from ....blackwell.custom_mix_cga_helpers import ( - TmaAtomOrPair, - pipeline_init_arrive_mixed_cga, - pipeline_init_wait_mixed_cga, -) -from ....blackwell.inference.mega.topk_reduce import TopkReduce from ....schedulers.base import WorkIdAcquisitionMode from ....schedulers.fc12_scheduler import ( BlackwellFusedFc12Scheduler, @@ -45,10 +38,12 @@ minimum_phase_interleave_hint, ) from ....schedulers.non_clc_mixed_cga import NonClcMixedCgaConfig +from ...custom_mix_cga_helpers import TmaAtomOrPair, pipeline_init_arrive_mixed_cga, pipeline_init_wait_mixed_cga from .block_scaled_swap_ab_fc12_epilogue import GatedActEpilogueArgs, SwapABGatedActEpilogue from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension from .block_scaled_swap_ab_fc12_mainloop import BlockScaledSwapAbFc12Mainloop + _aot_symbol_prefix = "rubin_mega_moe_aot" @@ -111,11 +106,24 @@ def impl_desc_require(cls) -> dict[str, type]: "fallback_cluster_count": OptionalRequirement(Optional[int]), "token_in_flag_batch": int, "token_back_mode": str, + "comm_clusters": OptionalRequirement(int), + "token_back_on_compute_ctas": OptionalRequirement(bool), + "comm_warps": OptionalRequirement(tuple), + "comm_stage_depths": OptionalRequirement(tuple), + "token_back_chunk_bytes": OptionalRequirement(int), + "dispatch_sf_group": OptionalRequirement(int), + "token_back_schedule_mode": OptionalRequirement(str), + "token_back_ready_granularity": OptionalRequirement(str), "reduce_topk_in_kernel": bool, + "helper_expert_count": OptionalRequirement(int), + "total_helper_slots": OptionalRequirement(int), } def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self._validate_desc_inputs(problem_desc, impl_desc) + problem_desc, impl_desc = resolve_dlb_expert_counts( + problem_desc, impl_desc, world_size=problem_desc["world_size"] + ) self.expert_count = problem_desc["expert_count"] self.intermediate_gateup_size = problem_desc["intermediate_gateup_size"] @@ -157,11 +165,47 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.launch_cluster_count = impl_desc["launch_cluster_count"] self.token_in_flag_batch = impl_desc["token_in_flag_batch"] self.token_back_mode = impl_desc["token_back_mode"] + self.token_back_ready_granularity = impl_desc.get("token_back_ready_granularity", "expert") self.reduce_topk_in_kernel = impl_desc["reduce_topk_in_kernel"] + self.comm_clusters = impl_desc.get("comm_clusters", 0) + self.token_back_on_compute_ctas = impl_desc.get("token_back_on_compute_ctas", False) + self.compute_cluster_count = self.launch_cluster_count - self.comm_clusters + # Left unbound rather than None when absent: the framework skips an optional field that was never + # bound but type-checks one that was, and only a split reads this, so there is no fused tuple to + # fall back to. + if impl_desc.get("comm_warps") is not None: + self.comm_warps = tuple(impl_desc["comm_warps"]) + self.comm_stage_depths = tuple(impl_desc.get("comm_stage_depths", (1, 1))) + # The fused default follows token_comm's own rather than restating it, so the two cannot drift. + self.token_back_chunk_bytes = impl_desc.get( + "token_back_chunk_bytes", TokenCommNonDeterministic.standalone_chunk_bytes + ) + self.dispatch_sf_group = impl_desc.get("dispatch_sf_group", 1) + # Dedicated communication CTAs use the static staged loop by default. + # Fused kernels retain the original work-ID-dependent token-back schedule. + default_schedule = ( + "static" if self.comm_clusters > 0 or self.work_id_mode != "atomic_counter" else "atomic_counter" + ) + self.token_back_schedule_mode = impl_desc.get("token_back_schedule_mode", default_schedule) + + self.helper_expert_count = impl_desc.get("helper_expert_count", 0) + if type(self.helper_expert_count) is not int or self.helper_expert_count < 0: + raise ValueError("helper_expert_count must be a non-negative exact int") self.occupancy = 1 self.architecture = "sm_107" self.local_expert_count = self.expert_count // self.world_size + if self.helper_expert_count > 0: + helper_topology = derive_helper_weight_ready_topology( + self.local_expert_count, + self.helper_expert_count, + ) + self.home_expert_count = helper_topology.home_expert_count + else: + # Preserve the baseline's runtime-expert-count specialization. A + # static M is required only when an S>0 helper boundary must be + # compiled into the scheduler gate. + self.home_expert_count = self.local_expert_count self.threads_per_cta = 16 * 32 if self.token_back_mode == "standalone_warps" else 12 * 32 # nosec B105 self.other_warp_register_count = 64 if self.token_back_mode == "standalone_warps" else 72 # nosec B105 @@ -176,9 +220,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.preferred_cluster_count = self.mixed_cga_config.preferred_cluster_count self.fallback_cluster_count = self.mixed_cga_config.fallback_cluster_count self.resolved_fallback_cluster_shape_mn = ( - self.cluster_shape_mn - if self.fallback_cluster_shape_mn is None - else self.fallback_cluster_shape_mn + self.cluster_shape_mn if self.fallback_cluster_shape_mn is None else self.fallback_cluster_shape_mn ) if self.launch_cluster_count != self.mixed_cga_config.launch_cluster_cnt_merge_as_preferred: raise ValueError( @@ -187,6 +229,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: ) self._validate_geometry() + self._validate_cta_partition() self._resolve_schedule_hint() mma_cta_count = 2 if self.use_2cta_instrs else 1 mma_instruction = Tcgen05MmaInstruction( @@ -200,32 +243,37 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: sf_vec_size=self.sf_vec_size, ) tmem_plan = make_tcgen05_tmem_plan(mma_instruction, self.architecture, self.mma_tiler_mnk) - tokens_per_fc1_ready_slot = self.mma_tiler_mnk[1] * self.cluster_shape_mn[1] - instruction_cta_m = self.mma_instruction_mnk[0] // mma_cta_count - instruction_m_repetitions = self.mma_tiler_mnk[0] // self.mma_instruction_mnk[0] - hidden_per_fc2_cluster_tile = ( - instruction_cta_m * instruction_m_repetitions * self.cluster_shape_mn[0] - ) - fc2_done_signals_per_token_tile = ( - ceil_div(self.hidden_size, hidden_per_fc2_cluster_tile) - * self.cluster_shape_mn[0] - * self.cluster_shape_mn[1] - ) + if self.token_back_ready_granularity == "token_tile": # nosec B105 + tokens_per_fc1_ready_slot, fc2_done_signals_per_token_tile = derive_fc2_ready_geometry( + hidden_size=self.hidden_size, + mma_tiler_mnk=self.mma_tiler_mnk, + cluster_shape_mn=self.cluster_shape_mn, + use_2cta_instrs=self.use_2cta_instrs, + fallback_cluster_shape_mn=self.fallback_cluster_shape_mn, + ) + else: + tokens_per_fc1_ready_slot = self.mma_tiler_mnk[1] * self.cluster_shape_mn[1] + instruction_cta_m = self.mma_instruction_mnk[0] // mma_cta_count + instruction_m_repetitions = self.mma_tiler_mnk[0] // self.mma_instruction_mnk[0] + hidden_per_fc2_cluster_tile = instruction_cta_m * instruction_m_repetitions * self.cluster_shape_mn[0] + fc2_done_signals_per_token_tile = ( + ceil_div(self.hidden_size, hidden_per_fc2_cluster_tile) + * self.cluster_shape_mn[0] + * self.cluster_shape_mn[1] + ) token_comm_impl_desc = ImplDesc( { **impl_desc, + **self._transfer_impl_desc_fields(), "fallback_cluster_shape_mn": self.fallback_cluster_shape_mn, "preferred_cluster_count": self.preferred_cluster_count, "fallback_cluster_count": self.fallback_cluster_count, + "token_back_ready_granularity": self.token_back_ready_granularity, "tokens_per_fc1_ready_slot": tokens_per_fc1_ready_slot, "fc2_done_signals_per_token_tile": fc2_done_signals_per_token_tile, "promised_launchable_sm_count": self.mixed_cga_config.total_cta_cnt, - "token_back_schedule_mode": ( - "atomic_counter" if self.work_id_mode == "atomic_counter" else "static" - ), - "router_smem_limit_bytes": cutlass.memory.get_smem_capacity_in_bytes( - self.architecture - ), + "token_back_schedule_mode": self.token_back_schedule_mode, + "router_smem_limit_bytes": utils.get_smem_capacity_in_bytes(self.architecture), } ) self.token_comm = TokenCommNonDeterministic(problem_desc, token_comm_impl_desc) @@ -248,15 +296,17 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: "token_back_push_data": self.token_comm.token_back_push_data, "fc1_epi_flag_batch": self.epi_flag_batches[0], "fc2_epi_flag_batch": self.epi_flag_batches[1], + "compute_cluster_count": self.compute_cluster_count, + "token_back_ready_granularity": self.token_back_ready_granularity, } ) scheduler_type = ( - BlackwellFusedFc12Scheduler - if self.schedule_mode == "grouped" - else PhaseInterleavedFc12Scheduler + BlackwellFusedFc12Scheduler if self.schedule_mode == "grouped" else PhaseInterleavedFc12Scheduler ) self.scheduler = scheduler_type(fc12_problem_desc, resolved_impl_desc) self.epilogue = SwapABGatedActEpilogue(fc12_problem_desc, resolved_impl_desc) + if self.token_back_ready_granularity == "token_tile": # nosec B105 + self._validate_token_tile_ready_contract() if self.epilogue.fc1_output_dtype is not self.b_dtype: raise ValueError("Epilogue FC1 output dtype must match the Mainloop B dtype.") if self.epilogue.fc1_output_sf_dtype is not self.sf_dtype: @@ -264,16 +314,150 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: if self.epilogue.sf_vec_size != self.sf_vec_size: raise ValueError("Epilogue and Mainloop scale vector sizes must match.") self._device_workspace = self._build_device_workspace() - self._mainloop, self._smem_workspace = self._build_mainloop_and_smem( - fc12_problem_desc, resolved_impl_desc - ) + self._comm_cta_mainloop = None + self._comm_smem_workspace = None + if self.comm_clusters > 0: + self._comm_cta_mainloop, self._comm_smem_workspace = self._build_comm_cta_and_smem() + self._mainloop, self._smem_workspace = self._build_mainloop_and_smem(fc12_problem_desc, resolved_impl_desc) + # Every CTA allocates the same arena and the two role tables are laid over it, each CTA reading only + # the one for the role it took. The arena therefore has to be the larger role's, and the workspace + # that owns it is whichever that is. + self._arena_workspace = self._smem_workspace + if ( + self._comm_smem_workspace is not None + and self._comm_smem_workspace.total_bytes > self._smem_workspace.total_bytes + ): + self._arena_workspace = self._comm_smem_workspace + self.smem_arena_bytes = self._arena_workspace.total_bytes + smem_limit = utils.get_smem_capacity_in_bytes(self.architecture) // self.occupancy + if self.smem_arena_bytes > smem_limit: + comm_bytes = 0 if self._comm_smem_workspace is None else self._comm_smem_workspace.total_bytes + raise ValueError( + f"CTA shared-memory arena requires {self.smem_arena_bytes} bytes " + f"(compute={self._smem_workspace.total_bytes}, comm={comm_bytes}), exceeding the " + f"{smem_limit}-byte limit for {self.architecture} at occupancy {self.occupancy}." + ) + if self._comm_smem_workspace is not None: + self._comm_smem_workspace.assert_fits_arena(self.smem_arena_bytes) + self._smem_workspace.assert_fits_arena(self.smem_arena_bytes) self._topk_reduce = ( - None - if self.reduce_topk_in_kernel - else TopkReduce(self.hidden_size, self.topk, self.combine_format) + None if self.reduce_topk_in_kernel else TopkReduce(self.hidden_size, self.topk, self.combine_format) ) + def _validate_cta_partition(self) -> None: + for name, value in (("launch_cluster_count", self.launch_cluster_count), ("comm_clusters", self.comm_clusters)): + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"{name} must be a Python int, got {value!r}.") + if not 0 <= self.comm_clusters < self.launch_cluster_count: + raise ValueError("comm_clusters must leave at least one compute cluster.") + if len(self.comm_stage_depths) != 2 or any( + isinstance(depth, bool) or not isinstance(depth, int) or not 1 <= depth <= 32 + for depth in self.comm_stage_depths + ): + raise ValueError("comm_stage_depths must be two ints in [1, 32], one phase bit per stage.") + if self.token_back_chunk_bytes <= 0 or self.token_back_chunk_bytes % 16: + raise ValueError("token_back_chunk_bytes must be a positive multiple of 16.") + if self.dispatch_sf_group not in (1, 4): + raise ValueError("dispatch_sf_group must be 1 or 4.") + if self.comm_clusters == 0: + if self.token_back_on_compute_ctas: + raise ValueError("token_back_on_compute_ctas requires a nonzero comm_clusters split.") + return + warps = getattr(self, "comm_warps", None) + if ( + warps is None + or len(warps) != 2 + or any(isinstance(w, bool) or not isinstance(w, int) or w <= 0 for w in warps) + ): + raise ValueError(f"A split needs comm_warps as two positive ints, got {warps}.") + if self.token_back_mode != "standalone_warps": # nosec B105 + raise ValueError("A CTA split requires token_back_mode='standalone_warps'.") + dispatch_warps, back_warps = warps + # Check the warps that actually coexist in each concrete CTA branch. + required_warps = ( + max(dispatch_warps, self.transfer_warp_idx_start + back_warps) + if self.token_back_on_compute_ctas + else dispatch_warps + back_warps + ) + if required_warps > self.threads_per_cta // 32: + raise ValueError( + f"comm_warps={warps} requires {required_warps} warps per CTA, " + f"but the launch provides only {self.threads_per_cta // 32}." + ) + dispatch_depth = self.comm_stage_depths[0] + if dispatch_depth > 1 and (self.dispatch_sf_group != 4 or dispatch_depth % 4): + raise ValueError("Multistage token-in needs dispatch_sf_group=4 and a depth divisible by four.") + + def _transfer_impl_desc_fields(self) -> dict: + return { + "comm_cta_count": self.comm_clusters * self.cluster_shape_mn[0] * self.cluster_shape_mn[1], + "comm_stage_depths": self.comm_stage_depths, + "token_back_chunk_bytes": self.token_back_chunk_bytes, + "dispatch_sf_group": self.dispatch_sf_group, + } + + def _is_dispatch_cta(self, cluster_idx): + # Partition preferred-shaped slots, keeping each actual MMA cluster intact. + return (cluster_idx * self.comm_clusters) % self.launch_cluster_count < self.comm_clusters + + def _compute_cluster_ordinal(self, cluster_idx): + preceding_dispatch = ( + cluster_idx * self.comm_clusters + self.launch_cluster_count - 1 + ) // self.launch_cluster_count + return cluster_idx - preceding_dispatch + + def _dispatch_cta_ordinal(self, cluster_idx, physical_cta_offset): + # An active fallback cluster rank restarts within a preferred slot; blockIdx.x does not. + slot = (cluster_idx * self.comm_clusters) // self.launch_cluster_count + return slot * self.cluster_shape_mn[0] + physical_cta_offset + + def _combine_cta_ordinal(self, cluster_idx, physical_cta_offset): + if self.token_back_on_compute_ctas: + return self._compute_cluster_ordinal(cluster_idx) * self.cluster_shape_mn[0] + physical_cta_offset + return self._dispatch_cta_ordinal(cluster_idx, physical_cta_offset) + + def _validate_token_tile_ready_contract(self) -> None: + """Check the consumer target against the actual scheduler and publisher. + + Both preferred clusters and a complete fallback group cover the same + logical CTA coordinates. Using active-cluster dimensions here would + permit token-back before the other fallback clusters finish their data. + """ + if not self.epilogue.fc2_use_tma: + raise ValueError("token_tile return-ready requires FC2 TMA stores.") + mapping_tile = self.scheduler.mapping_cta_tile_shape_mnk + mapping_cluster = self.scheduler.mapping_cluster_shape_mn + if ( + mapping_cluster != (1, self.cluster_shape_mn[0]) + or mapping_tile[:2] != (self.epilogue.cta_tile_n, self.epilogue.cta_tile_m) + ): + raise ValueError("token_tile scheduler and epilogue must share the preferred logical CTA geometry.") + expected_signals = ceil_div(self.hidden_size, mapping_tile[1] * mapping_cluster[1]) * mapping_cluster[1] + if ( + self.token_comm.tokens_per_fc1_ready_slot != mapping_tile[0] + or self.token_comm.fc2_done_signals_per_token_tile != expected_signals + ): + raise ValueError("token_tile producer and consumer must share logical ready slots and completion counts.") + def _validate_geometry(self) -> None: + if self.token_back_ready_granularity not in ("expert", "token_tile"): + raise ValueError(f"Unsupported token_back_ready_granularity {self.token_back_ready_granularity!r}.") + if self.token_back_ready_granularity == "token_tile": # nosec B105 + if not ( + self.schedule_mode in ("grouped", "phase_interleave") + and self.work_id_mode == "atomic_counter" + and self.token_back_mode in ("standalone_warps", "reuse_dispatch_warps") + and self.use_2cta_instrs + and self.fc2_use_bulk + and self.quant_kind == QuantKind.nvfp4 + and not self.reduce_topk_in_kernel + ): + raise ValueError( + "token_tile return-ready requires grouped or phase_interleave/atomic_counter, " + "standalone_warps or reuse_dispatch_warps, " + "two-CTA MMA with hidden-only preferred/fallback clusters, " + "FC2 TMA stores, NVFP4, and separate top-k reduction." + ) static_expert_dimensions = ( isinstance(self.expert_count, int), isinstance(self.intermediate_gateup_size, int), @@ -282,9 +466,7 @@ def _validate_geometry(self) -> None: if any(static_expert_dimensions) and not all(static_expert_dimensions): raise ValueError("MegaMoE expert dimensions must be either all static or all runtime.") if not all(static_expert_dimensions): - raise NotImplementedError( - "The MegaMoE Kernel currently requires static expert dimensions." - ) + raise NotImplementedError("The MegaMoE Kernel currently requires static expert dimensions.") if self.expert_count <= 0 or self.intermediate_gateup_size <= 0 or self.hidden_size <= 0: raise ValueError("MegaMoE expert dimensions must be positive.") if self.world_size <= 0 or self.expert_count % self.world_size != 0: @@ -292,9 +474,7 @@ def _validate_geometry(self) -> None: if self.topk <= 0 or self.topk > self.expert_count: raise ValueError("topk must be positive and no greater than expert_count.") if self.topk_index_dtype not in (cutlass.Int32, cutlass.Int64): - raise ValueError( - f"topk_index_dtype must be Int32 or Int64, got {self.topk_index_dtype}." - ) + raise ValueError(f"topk_index_dtype must be Int32 or Int64, got {self.topk_index_dtype}.") if self.reduce_topk_in_kernel and not self.apply_topk_at_fc1: raise ValueError("In-kernel top-k reduction requires apply_topk_at_fc1=True.") if self.intermediate_gateup_size % 2 != 0: @@ -312,12 +492,8 @@ def _validate_geometry(self) -> None: raise ValueError( f"schedule_policy mode must be 'grouped' or 'phase_interleave', got {self.schedule_mode!r}." ) - if self.hint is not None and ( - isinstance(self.hint, bool) or not isinstance(self.hint, int) or self.hint <= 0 - ): - raise ValueError( - f"schedule_policy hint must be a positive Python int, got {self.hint!r}." - ) + if self.hint is not None and (isinstance(self.hint, bool) or not isinstance(self.hint, int) or self.hint <= 0): + raise ValueError(f"schedule_policy hint must be a positive Python int, got {self.hint!r}.") if self.schedule_mode == "phase_interleave" and self.work_id_mode != "atomic_counter": raise ValueError("phase_interleave currently requires work_id_mode='atomic_counter'.") if len(self.epi_flag_batches) != 2: @@ -335,20 +511,14 @@ def _validate_geometry(self) -> None: ): if len(dimensions) != expected_rank: raise ValueError(f"{field_name} must contain {expected_rank} dimensions.") - if not all( - isinstance(dimension, int) and not isinstance(dimension, bool) - for dimension in dimensions - ): + if not all(isinstance(dimension, int) and not isinstance(dimension, bool) for dimension in dimensions): raise TypeError(f"{field_name} dimensions must be Python integers.") if any(dimension <= 0 for dimension in dimensions): raise ValueError(f"{field_name} dimensions must be positive.") if self.quant_kind == QuantKind.mxfp4_mxfp8: # Rubin consumes the mixed FP4 weight from native packed SMEM. Its K-major source only # needs ordinary 16-byte TMA alignment and complete SFVec32 blocks. - for name, extent in ( - ("hidden_size", self.hidden_size), - ("intermediate_downproj", intermediate_downproj), - ): + for name, extent in (("hidden_size", self.hidden_size), ("intermediate_downproj", intermediate_downproj)): if extent * int(self.a_dtype.width) % 128 != 0 or extent % self.sf_vec_size != 0: raise ValueError( f"Rubin {self.quant_kind} requires {name} to satisfy ordinary TMA and " @@ -372,9 +542,7 @@ def _validate_geometry(self) -> None: f"{self.quant_kind} requires Rubin 2x instruction K={expected_instruction_k}, got {instruction_k}." ) if (instruction_m, instruction_n) != (mma_m, mma_n): - raise NotImplementedError( - "Rubin MegaMoE does not implement M/N instruction repetition or B-reuse." - ) + raise NotImplementedError("Rubin MegaMoE does not implement M/N instruction repetition or B-reuse.") if mma_n not in (64, 128, 256): raise ValueError(f"mma_tiler N must be 64, 128, or 256, got {mma_n}.") if mma_k % instruction_k != 0: @@ -393,9 +561,7 @@ def _validate_geometry(self) -> None: f"four-vector SF atom ({sf_atom_size} elements)." ) elif self.intermediate_gateup_size % (self.sf_vec_size * 4) != 0: - raise ValueError( - "The intermediate dimension must be divisible by four scale-factor vectors." - ) + raise ValueError("The intermediate dimension must be divisible by four scale-factor vectors.") if cluster_n != 1: raise ValueError(f"The swap-AB FC12 path requires cluster N=1, got {cluster_n}.") if cluster_m <= 0 or cluster_m > 16 or cluster_m & (cluster_m - 1): @@ -405,14 +571,8 @@ def _validate_geometry(self) -> None: if self.fallback_cluster_shape_mn is not None: fallback_cluster_m, fallback_cluster_n = self.fallback_cluster_shape_mn if fallback_cluster_n != 1: - raise ValueError( - f"The swap-AB FC12 fallback path requires cluster N=1, got {fallback_cluster_n}." - ) - if ( - fallback_cluster_m <= 0 - or fallback_cluster_m > 16 - or fallback_cluster_m & (fallback_cluster_m - 1) - ): + raise ValueError(f"The swap-AB FC12 fallback path requires cluster N=1, got {fallback_cluster_n}.") + if fallback_cluster_m <= 0 or fallback_cluster_m > 16 or fallback_cluster_m & (fallback_cluster_m - 1): raise ValueError("fallback cluster M must be a power of two no greater than 16.") if self.use_2cta_instrs and fallback_cluster_m % 2 != 0: raise ValueError("Two-CTA MMA requires an even fallback cluster M.") @@ -428,7 +588,7 @@ def _resolve_schedule_hint(self) -> None: raw_minimum_hint = minimum_phase_interleave_hint( blocks_fc1=blocks_fc1, blocks_fc2=blocks_fc2, - launch_cluster_cnt_merge_as_preferred=self.mixed_cga_config.launch_cluster_cnt_merge_as_preferred, + launch_cluster_cnt_merge_as_preferred=self.compute_cluster_count, ) minimum_safe_hint = round_up(raw_minimum_hint, self.epi_flag_batches[0]) if self.hint is None: @@ -508,6 +668,7 @@ def dtype_name(dtype: type) -> str: f"{dtype_name(self.a_dtype)}_{dtype_name(self.b_dtype)}_{dtype_name(self.acc_dtype)}_" f"sfvec{self.sf_vec_size}_a{self.a_major_mode.name.lower()}_b{self.b_major_mode.name.lower()}_" f"e{self.expert_count}_ep{self.world_size}_topk{self.topk}_" + f"helpers{self.helper_expert_count}_" f"topkidx{dtype_name(self.topk_index_dtype)}_" f"h{self.hidden_size}_i{self.intermediate_gateup_size}_maxtoken{self.max_tokens_per_rank}_" f"inst{instruction}_{self.mma_k_mode}_tile{tile}_" @@ -526,6 +687,8 @@ def dtype_name(dtype: type) -> str: f"situ{self.situ_beta}x{self.situ_linear_beta}_" f"{'apply_topk_fc1' if self.apply_topk_at_fc1 else 'apply_topk_fc2'}_" f"{'inkernel_redg' if self.reduce_topk_in_kernel else 'separate_reduce'}" + "_dlb_v1" + + ("_fc2readytile_v1" if self.token_back_ready_granularity == "token_tile" else "") # nosec B105 ) def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): @@ -547,14 +710,10 @@ def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): extents = tuple( - sym_int64(divisibility=math.gcd(int(extent), 128)) - if axis in dynamic_axes - else int(extent) + sym_int64(divisibility=math.gcd(int(extent), 128)) if axis in dynamic_axes else int(extent) for axis, extent in enumerate(shape) ) - return make_fake_compact_tensor( - dtype, extents, stride_order=stride_order, assumed_align=alignment - ) + return make_fake_compact_tensor(dtype, extents, stride_order=stride_order, assumed_align=alignment) tokens = self.max_tokens_per_rank hidden = self.hidden_size @@ -563,20 +722,14 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): experts = self.local_expert_count sf_vec_size = self.sf_vec_size # Weight SF is atom-swizzled and opaque, so its extent is whatever the caller staged, not a kernel choice. - fc1_weight_sf_columns = round_up(intermediate_gateup, 128) * round_up( - hidden // sf_vec_size, 4 - ) - fc2_weight_sf_columns = round_up(hidden, 128) * round_up( - intermediate_downproj // sf_vec_size, 4 - ) + fc1_weight_sf_columns = round_up(intermediate_gateup, 128) * round_up(hidden // sf_vec_size, 4) + fc2_weight_sf_columns = round_up(hidden, 128) * round_up(intermediate_downproj // sf_vec_size, 4) # TopkReduce and the in-kernel REDG both take their output dtype from this tensor. output_dtype = cutlass.BFloat16 fake_arguments = dict( # Swap-AB puts the weights on the MMA A operand, so the token wire dtype is TokenComm's, not a_dtype. - activation=fake_tensor( - self.token_comm.activation_dtype, (tokens, hidden), (1, 0), {0}, 16 - ), + activation=fake_tensor(self.token_comm.activation_dtype, (tokens, hidden), (1, 0), {0}, 16), activation_sf=fake_tensor( self.token_comm.activation_sf_dtype, (tokens, self.token_comm.activation_sf_hidden_padded), @@ -586,18 +739,10 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): ), topk_indices=fake_tensor(self.topk_index_dtype, (tokens, self.topk), (1, 0), {0}, 16), topk_scores=fake_tensor(cutlass.Float32, (tokens, self.topk), (1, 0), {0}, 4), - fc1_weight=fake_tensor( - self.a_dtype, (experts, hidden, intermediate_gateup), (2, 0, 1), {0, 2}, 16 - ), - fc1_weight_sf=fake_tensor( - self.sf_dtype, (experts, fc1_weight_sf_columns), (1, 0), {0}, 16 - ), - fc2_weight=fake_tensor( - self.a_dtype, (experts, intermediate_downproj, hidden), (2, 0, 1), {0, 2}, 16 - ), - fc2_weight_sf=fake_tensor( - self.sf_dtype, (experts, fc2_weight_sf_columns), (1, 0), {0}, 16 - ), + fc1_weight=fake_tensor(self.a_dtype, (experts, hidden, intermediate_gateup), (2, 0, 1), {0, 2}, 16), + fc1_weight_sf=fake_tensor(self.sf_dtype, (experts, fc1_weight_sf_columns), (1, 0), {0}, 16), + fc2_weight=fake_tensor(self.a_dtype, (experts, intermediate_downproj, hidden), (2, 0, 1), {0, 2}, 16), + fc2_weight_sf=fake_tensor(self.sf_dtype, (experts, fc2_weight_sf_columns), (1, 0), {0}, 16), output_activation=fake_tensor(output_dtype, (tokens, hidden), (1, 0), {0}, 16), # Opaque byte workspaces: a bare base pointer, like to_cute_ptr. The 128-byte promise matches the # alignment the registered regions assume, so the caller must hand over 128-byte-aligned bases. @@ -612,6 +757,16 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): ), stream=make_fake_stream(), ) + if self.helper_expert_count > 0: + fake_arguments.update( + hot_expert_weight_ready_flags=make_ptr( + cutlass.Uint64, + 0, + AddressSpace.gmem, + assumed_align=8, + ), + hot_expert_weight_ready_generation=cutlass.Uint64(1), + ) if self.quant_kind.uses_global_scale: # nvfp4 carries per-expert dequant scalars. Under an e8m0 scale they do not exist at # all, and declaring them here would put three tensors in the ABI that the caller has @@ -626,22 +781,40 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): if out_path is None: return compiled compiled.export_to_c( - out_path, function_name=_aot_symbol_prefix, export_only_tvm_ffi_symbols=True + out_path, + function_name=self._aot_symbol_name(), + export_only_tvm_ffi_symbols=True, ) return out_path - @staticmethod - def load_compiled(path: str): + def _aot_symbol_name(self) -> str: + identity = hashlib.sha256(self.name().encode()).hexdigest()[:16] + return f"{_aot_symbol_prefix}_{identity}" + + def load_compiled(self, path: str): from cutlass.cute.runtime import load_module - return load_module(path, enable_tvm_ffi=True)[_aot_symbol_prefix] + function = getattr(load_module(path, enable_tvm_ffi=True), self._aot_symbol_name()) + return make_megamoe_aot_callable(function, self) + + def _build_comm_cta_and_smem(self) -> Tuple[TokenCommCtaMainloop, SmemWorkspace]: + mainloop = TokenCommCtaMainloop( + self.token_comm, token_back_on_compute_ctas=self.token_back_on_compute_ctas + ) + workspace = SmemWorkspace() + mainloop.register_smem_regions(workspace) + workspace.finalize() + return mainloop, workspace def _build_mainloop_and_smem( self, problem_desc: ProblemDesc, resolved_impl_desc: ImplDesc ) -> Tuple[BlockScaledSwapAbFc12Mainloop, SmemWorkspace]: - smem_limit = cutlass.memory.get_smem_capacity_in_bytes(self.architecture) // self.occupancy + smem_limit = utils.get_smem_capacity_in_bytes(self.architecture) // self.occupancy smem_workspace = SmemWorkspace() - self.token_comm.register_smem_regions(smem_workspace) + if self.comm_clusters == 0: + self.token_comm.register_smem_regions(smem_workspace) + elif self.token_back_on_compute_ctas: + self._comm_cta_mainloop.register_token_back_smem(smem_workspace) self.scheduler.register_smem_regions(smem_workspace) self.epilogue.register_smem_regions(smem_workspace) fixed_component_bytes = smem_workspace.estimate_total_bytes() @@ -650,15 +823,14 @@ def _build_mainloop_and_smem( if mainloop_smem_budget_bytes <= 0: raise ValueError("Fixed kernel components leave no SMEM budget for the Mainloop.") - mainloop_impl_desc = ImplDesc( - {**resolved_impl_desc, "mainloop_smem_budget_bytes": mainloop_smem_budget_bytes} - ) + mainloop_impl_desc = ImplDesc({**resolved_impl_desc, "mainloop_smem_budget_bytes": mainloop_smem_budget_bytes}) mainloop = BlockScaledSwapAbFc12Mainloop(problem_desc, mainloop_impl_desc) mainloop.register_smem_regions(smem_workspace) smem_workspace.finalize(max_bytes=smem_limit) return mainloop, smem_workspace + @cute.jit def __call__( self, @@ -678,16 +850,34 @@ def __call__( fc1_alpha: Optional[cute.Tensor] = None, # (local_experts,) fc2_alpha: Optional[cute.Tensor] = None, # (local_experts,) fc1_norm_const: Optional[cute.Tensor] = None, # (local_experts,) + hot_expert_weight_ready_flags: Optional[cute.Pointer] = None, # uint64[EP] + hot_expert_weight_ready_generation: Optional[cutlass.Uint64] = None, ) -> None: """Launch router, fused MegaMoE compute, and optional TopK reduce.""" + helper_ready_bundle = ( + hot_expert_weight_ready_flags, + hot_expert_weight_ready_generation, + ) + if cutlass.const_expr( + ( + self.helper_expert_count > 0 + and not all(value is not None for value in helper_ready_bundle) + ) + or ( + self.helper_expert_count == 0 + and any(value is not None for value in helper_ready_bundle) + ) + ): + raise ValueError( + "helper READY flags/generation must both be present iff helper_expert_count > 0" + ) + def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: return cute.make_tensor(tensor.iterator, cute.make_layout(shape, stride=tensor.stride)) if cutlass.const_expr(topk_indices.element_type is not self.topk_index_dtype): - raise TypeError( - f"topk_indices must be the declared {self.topk_index_dtype}, got {topk_indices.dtype}." - ) + raise TypeError(f"topk_indices must be the declared {self.topk_index_dtype}, got {topk_indices.dtype}.") for operand_name, operand, expected_dtype in ( ("activation", activation, self.b_dtype), ("activation_sf", activation_sf, self.sf_dtype), @@ -776,19 +966,14 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: padded_intermediate_gateup = cute.round_up(intermediate_gateup, sf_vec_size * 4) expected_fc1_weight_sf_columns = padded_intermediate_gateup * padded_hidden // sf_vec_size if cutlass.const_expr( - isinstance(fc1_weight_sf.shape[1], int) - and isinstance(expected_fc1_weight_sf_columns, int) + isinstance(fc1_weight_sf.shape[1], int) and isinstance(expected_fc1_weight_sf_columns, int) ): if cutlass.const_expr(fc1_weight_sf.shape[1] != expected_fc1_weight_sf_columns): raise ValueError("fc1_weight_sf has an incompatible column count.") fc1_sfa = cute.make_tensor( fc1_weight_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - cutlass.Int32(padded_intermediate_gateup), - cutlass.Int32(padded_hidden), - cutlass.Int32(experts), - ), + (cutlass.Int32(padded_intermediate_gateup), cutlass.Int32(padded_hidden), cutlass.Int32(experts)), sf_vec_size, ), ) @@ -806,31 +991,21 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: padded_fc2_intermediate = cute.round_up(intermediate_fc2, sf_vec_size * 4) expected_fc2_weight_sf_columns = padded_fc2_hidden * padded_fc2_intermediate // sf_vec_size if cutlass.const_expr( - isinstance(fc2_weight_sf.shape[1], int) - and isinstance(expected_fc2_weight_sf_columns, int) + isinstance(fc2_weight_sf.shape[1], int) and isinstance(expected_fc2_weight_sf_columns, int) ): if cutlass.const_expr(fc2_weight_sf.shape[1] != expected_fc2_weight_sf_columns): raise ValueError("fc2_weight_sf has an incompatible column count.") fc2_sfa = cute.make_tensor( fc2_weight_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - cutlass.Int32(padded_fc2_hidden), - cutlass.Int32(padded_fc2_intermediate), - cutlass.Int32(experts_fc2), - ), + (cutlass.Int32(padded_fc2_hidden), cutlass.Int32(padded_fc2_intermediate), cutlass.Int32(experts_fc2)), sf_vec_size, ), ) fc2_sfb = cute.make_tensor( fc1_output_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - fc1_output_sf.shape[0], - cutlass.Int32(fc1_output_sf.shape[1] * sf_vec_size), - singleton, - ), - sf_vec_size, + (fc1_output_sf.shape[0], cutlass.Int32(fc1_output_sf.shape[1] * sf_vec_size), singleton), sf_vec_size ), ) if cutlass.const_expr(self.reduce_topk_in_kernel): @@ -838,28 +1013,19 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: output_activation.iterator, cute.make_layout( (self.max_tokens_per_rank, 1, hidden), - stride=( - output_activation.stride[0], - output_activation.stride[0], - output_activation.stride[1], - ), + stride=(output_activation.stride[0], output_activation.stride[0], output_activation.stride[1]), ), ) pre_reduced_activation_sf = None else: - pre_reduced_activation = self.token_comm.pre_reduced_activation_tensor( - self._device_workspace - ) - pre_reduced_activation_sf = self.token_comm.pre_reduced_activation_sf_tensor( - self._device_workspace - ) + pre_reduced_activation = self.token_comm.pre_reduced_activation_tensor(self._device_workspace) + pre_reduced_activation_sf = self.token_comm.pre_reduced_activation_sf_tensor(self._device_workspace) if cutlass.const_expr(self.token_comm.token_back_push_data): fc2_output = self.token_comm.fc2_activation_tensor(self._device_workspace) else: fc2_output = rewrite_tensor_shape( - pre_reduced_activation, - (pre_reduced_activation.shape[0], pre_reduced_activation.shape[1], hidden), + pre_reduced_activation, (pre_reduced_activation.shape[0], pre_reduced_activation.shape[1], hidden) ) fc2_output_sf = self.token_comm.fc2_activation_sf_tensor(self._device_workspace) @@ -898,7 +1064,11 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: grid = self.scheduler.get_grid_shape(max_active_clusters=self.launch_cluster_count) self._device_workspace.remove_device_members() - kernel = self._kernel( + if cutlass.const_expr(self.comm_clusters > 0): + kernel_entry = self._kernel_cta_specialized + else: + kernel_entry = self._kernel + kernel = kernel_entry( fc1_tma_a_tensor, fc1_tma_a_atom, fc1_tma_sfa_tensor, @@ -933,6 +1103,8 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: fc1_alpha, fc2_alpha, fc1_norm_const, + hot_expert_weight_ready_flags, + hot_expert_weight_ready_generation, ) if cutlass.const_expr(self.mixed_cga_config.is_mixed): kernel.launch( @@ -956,13 +1128,10 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: if cutlass.const_expr(not self.reduce_topk_in_kernel): reduce_scores = None if cutlass.const_expr(self.apply_topk_at_fc1) else topk_scores self._topk_reduce( - pre_reduced_activation, - pre_reduced_activation_sf, - output_activation, - reduce_scores, - stream, + pre_reduced_activation, pre_reduced_activation_sf, output_activation, reduce_scores, stream ) + @cute.kernel def _kernel( self, @@ -1000,6 +1169,8 @@ def _kernel( fc1_alpha: Optional[cute.Tensor], fc2_alpha: Optional[cute.Tensor], fc1_norm_const: Optional[cute.Tensor], + hot_expert_weight_ready_flags: Optional[cute.Pointer], + hot_expert_weight_ready_generation: Optional[cutlass.Uint64], ): """Compose TokenComm, Scheduler, Mainloop, and Epilogue.""" self._mainloop.materialize_codegen_members() @@ -1021,8 +1192,7 @@ def _kernel( fc1_output_sf = cute.make_tensor( fc1_output_sf_storage.iterator, blockscaled_utils.tile_atom_to_shape_SF( - (fc1_output_sf_storage.shape[0], intermediate_gateup // 2, cutlass.Int32(1)), - self.epilogue.sf_vec_size, + (fc1_output_sf_storage.shape[0], intermediate_gateup // 2, cutlass.Int32(1)), self.epilogue.sf_vec_size ), ) @@ -1033,9 +1203,7 @@ def _kernel( is_fallback_cluster = Boolean(False) if cutlass.const_expr(self.mixed_cga_config.is_mixed): active_cluster_m, _, _ = cute.arch.block_in_cluster_dim() - is_fallback_cluster = active_cluster_m == Int32( - self.resolved_fallback_cluster_shape_mn[0] - ) + is_fallback_cluster = active_cluster_m == Int32(self.resolved_fallback_cluster_shape_mn[0]) cta_rank_in_cluster = cute.arch.block_idx_in_cluster() cta_coord_in_cluster = (cta_rank_in_cluster, Int32(0), Int32(0)) linear_cta_idx = block_idx[0] + block_idx[2] * Int32(self.cluster_shape_mn[0]) @@ -1058,12 +1226,7 @@ def _kernel( pool_topk_scores = self.token_comm.fc1_topk_scores_tensor(self._device_workspace) self._mainloop.assign_device_members( - self._smem_workspace, - smem_base, - cta_coord_in_cluster, - is_fallback_cluster, - hidden, - intermediate_gateup, + self._smem_workspace, smem_base, cta_coord_in_cluster, is_fallback_cluster, hidden, intermediate_gateup ) ab_pipeline = self._mainloop.create_ab_pipeline(self._smem_workspace, smem_base) tma_a_pipeline_state = pipeline.make_pipeline_state( @@ -1091,39 +1254,34 @@ def _kernel( ) scheduler = self.scheduler - fc2_spin_threshold = ( - intermediate_gateup + self._mainloop.cta_tile_m - 1 - ) // self._mainloop.cta_tile_m + fc2_spin_threshold = (intermediate_gateup + self._mainloop.cta_tile_m - 1) // self._mainloop.cta_tile_m + helper_weight_ready_gate = None + if cutlass.const_expr(self.helper_expert_count > 0): + helper_weight_ready_gate = make_helper_weight_ready_generation_gate( + memory_slot_count=self.local_expert_count, + helper_count=self.helper_expert_count, + source_count=self.world_size, + terminal_flags=hot_expert_weight_ready_flags, + expected_generation=hot_expert_weight_ready_generation, + ) kernel_extension = BlockScaledSwapAbFc12Extension( sf_vec_size=self.sf_vec_size, fc1_done_counter_pointer=fc1_done_counter.iterator, fc2_spin_threshold=fc2_spin_threshold, - fc1_ready_counter_pointer=self.token_comm.fc1_ready_counter_pointer( - self._device_workspace - ), + fc1_ready_counter_pointer=self.token_comm.fc1_ready_counter_pointer(self._device_workspace), ) optional_epilogue_args = GatedActEpilogueArgs( - fc1_alpha=fc1_alpha, - fc2_alpha=fc2_alpha, - fc1_norm_const=fc1_norm_const, - topk_scores=pool_topk_scores, + fc1_alpha=fc1_alpha, fc2_alpha=fc2_alpha, fc1_norm_const=fc1_norm_const, topk_scores=pool_topk_scores ) pipeline_init_arrive_mixed_cga( - self.cluster_shape_mn, - self.resolved_fallback_cluster_shape_mn, - is_fallback_cluster, - is_relaxed=True, + self.cluster_shape_mn, self.resolved_fallback_cluster_shape_mn, is_fallback_cluster, is_relaxed=True ) pipeline_init_wait_mixed_cga( self.cluster_shape_mn, self.resolved_fallback_cluster_shape_mn, is_fallback_cluster ) - finalize_barrier_thread_count = ( - 12 * 32 if self.token_back_mode == "standalone_warps" else 8 * 32 # nosec B105 - ) - finalize_barrier = pipeline.NamedBarrier( - barrier_id=13, num_threads=finalize_barrier_thread_count - ) + finalize_barrier_thread_count = 12 * 32 if self.token_back_mode == "standalone_warps" else 8 * 32 # nosec B105 + finalize_barrier = pipeline.NamedBarrier(barrier_id=13, num_threads=finalize_barrier_thread_count) if warp_idx == self.scheduler_warp_id: cute.arch.warpgroup_reg_dealloc(self.scheduler_register_count) @@ -1138,9 +1296,15 @@ def _kernel( iket.range_push("scheduler.gen_work") work_tile = scheduler.gen_next_work() iket.range_pop() + helper_ready_latched = Int32(0) while work_tile.is_valid_tile: iket.range_push("scheduler.publish_work") - scheduler.publish_work(kernel_extension.prepare_work_tile(work_tile)) + prepared_tile = kernel_extension.prepare_work_tile(work_tile) + if cutlass.const_expr(helper_weight_ready_gate is not None): + helper_ready_latched = helper_weight_ready_gate.wait_before_first_helper_publish( + prepared_tile.expert_idx, helper_ready_latched + ) + scheduler.publish_work(prepared_tile) iket.range_pop() iket.range_push("scheduler.gen_work") work_tile = scheduler.gen_next_work() @@ -1235,9 +1399,7 @@ def _kernel( optional_epilogue_args, ) tmem_allocator.relinquish_alloc_permit() - tmem_allocator.free( - tmem_allocator.retrieve_ptr(self.acc_dtype), self._mainloop.num_tmem_alloc_cols - ) + tmem_allocator.free(tmem_allocator.retrieve_ptr(self.acc_dtype), self._mainloop.num_tmem_alloc_cols) iket.range_pop() finalize_barrier.arrive_unaligned() @@ -1248,8 +1410,7 @@ def _kernel( self.token_comm.token_in(self._smem_workspace, smem_base) iket.range_pop() if cutlass.const_expr( - self.token_comm.token_back_enabled - and self.token_back_mode != "standalone_warps" # nosec B105 + self.token_comm.token_back_enabled and self.token_back_mode != "standalone_warps" # nosec B105 ): iket.range_push("mega.token_back") self.token_comm.token_back(self._smem_workspace, smem_base) @@ -1265,3 +1426,301 @@ def _kernel( finalize_barrier.arrive_unaligned() self.token_comm.remove_device_members() self._device_workspace.remove_device_members() + + + @cute.kernel + def _kernel_cta_specialized( + self, + fc1_tma_a_tensor: cute.Tensor, + fc1_tma_a_atom: cute.CopyAtom, + fc1_tma_sfa_tensor: cute.Tensor, + fc1_tma_sfa_atom: cute.CopyAtom, + fc2_tma_a_tensor: cute.Tensor, + fc2_tma_a_atom: cute.CopyAtom, + fc2_tma_sfa_tensor: cute.Tensor, + fc2_tma_sfa_atom: cute.CopyAtom, + fc1_tma_b_tensor: cute.Tensor, + fc1_tma_b_atom_or_pair: TmaAtomOrPair, + fc1_tma_sfb_tensor: cute.Tensor, + fc1_tma_sfb_atom_or_pair: TmaAtomOrPair, + fc2_tma_b_tensor: cute.Tensor, + fc2_tma_b_atom_or_pair: TmaAtomOrPair, + fc2_tma_sfb_tensor: cute.Tensor, + fc2_tma_sfb_atom_or_pair: TmaAtomOrPair, + fc1_output_tma_atom: cute.CopyAtom, + fc1_output_tma_tensor: cute.Tensor, + fc2_output_tma_atom: Optional[cute.CopyAtom], + fc2_output_tma_tensor: Optional[cute.Tensor], + fc2_output: cute.Tensor, + fc2_output_sf: Optional[cute.Tensor], + actual_expert_shape: Tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32], + activation: cute.Tensor, + activation_sf: cute.Tensor, + pre_reduced_activation: cute.Tensor, + pre_reduced_activation_sf: Optional[cute.Tensor], + peer_rank_ptr_mapper, + local_rank: cutlass.Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + fc1_alpha: Optional[cute.Tensor], + fc2_alpha: Optional[cute.Tensor], + fc1_norm_const: Optional[cute.Tensor], + hot_expert_weight_ready_flags: Optional[cute.Pointer], + hot_expert_weight_ready_generation: Optional[cutlass.Uint64], + ): + """Run communication and FC12 on distinct, whole-cluster CTA sets.""" + storage_type = self._arena_workspace.storage_class() + smem_allocator = utils.SmemAllocator() + storage = smem_allocator.allocate(storage_type) + smem_base = storage.buffer.data_ptr() + workspace = self._device_workspace + workspace.assign_device_members(local_workspace, shared_workspace) + thread_idx, _, _ = cute.arch.thread_idx() + warp_idx = thread_idx // 32 + block_idx = cute.arch.block_idx() + active_cluster_m = Int32(self.cluster_shape_mn[0]) + is_fallback_cluster = Boolean(False) + if cutlass.const_expr(self.mixed_cga_config.is_mixed): + active_cluster_m, _, _ = cute.arch.block_in_cluster_dim() + is_fallback_cluster = active_cluster_m == Int32(self.resolved_fallback_cluster_shape_mn[0]) + cta_rank_in_cluster = cute.arch.block_idx_in_cluster() + cta_coord_in_cluster = (cta_rank_in_cluster, Int32(0), Int32(0)) + linear_cta_idx = block_idx[0] + block_idx[2] * Int32(self.cluster_shape_mn[0]) + token_comm_args = TokenCommArgs( + activation=activation, + activation_sf=activation_sf, + pre_reduced_activation=pre_reduced_activation, + pre_reduced_activation_sf=pre_reduced_activation_sf, + peer_rank_ptr_mapper=peer_rank_ptr_mapper, + ) + self.token_comm.assign_device_members( + device_workspace=workspace, + token_comm_args=token_comm_args, + local_rank=local_rank, + linear_cta_idx=linear_cta_idx, + dispatch_cta_ordinal=self._dispatch_cta_ordinal(block_idx[2], block_idx[0]), + combine_cta_ordinal=self._combine_cta_ordinal(block_idx[2], block_idx[0]), + ) + if self._is_dispatch_cta(block_idx[2]): + self._comm_cta_mainloop.run(self._comm_smem_workspace, smem_base) + else: + compute_block_idx = (block_idx[0], block_idx[1], self._compute_cluster_ordinal(block_idx[2])) + self._mainloop.materialize_codegen_members() + fc1_done_counter = self._device_workspace.tensor(self.fc1_done_counter_region) + fc1_output_sf_storage = self._device_workspace.tensor(self.fc1_output_sf_region) + if cutlass.const_expr(isinstance(self.hidden_size, int)): + hidden = self.hidden_size + else: + hidden = actual_expert_shape[2] + if cutlass.const_expr(isinstance(self.intermediate_gateup_size, int)): + intermediate_gateup = self.intermediate_gateup_size + else: + intermediate_gateup = actual_expert_shape[1] + fc1_output_sf = cute.make_tensor( + fc1_output_sf_storage.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (fc1_output_sf_storage.shape[0], intermediate_gateup // 2, cutlass.Int32(1)), self.epilogue.sf_vec_size + ), + ) + + expert_token_sizes = self.token_comm.local_expert_sizes(self._device_workspace, local_rank) + token_src_metadata = self.token_comm.token_src_metadata_tensor(self._device_workspace) + fc2_done_counter = self.token_comm.fc2_done_counter_tensor(self._device_workspace) + pool_topk_scores = self.token_comm.fc1_topk_scores_tensor(self._device_workspace) + + self._mainloop.assign_device_members( + self._smem_workspace, smem_base, cta_coord_in_cluster, is_fallback_cluster, hidden, intermediate_gateup + ) + ab_pipeline = self._mainloop.create_ab_pipeline(self._smem_workspace, smem_base) + tma_a_pipeline_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self._mainloop.num_ab_pipeline_stages + ) + tma_b_pipeline_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self._mainloop.num_ab_pipeline_stages + ) + mma_pipeline_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self._mainloop.num_ab_pipeline_stages + ) + acc_pipeline = self._mainloop.create_acc_pipeline(self._smem_workspace, smem_base) + tmem_allocator = self._mainloop.create_tmem_allocator( + self._smem_workspace, smem_base, allocator_warp_id=self.epilogue_warp_ids[0] + ) + self.scheduler.assign_device_members( + expert_token_sizes=expert_token_sizes, + expert_token_prefix_sum=None, + actual_expert_shape=actual_expert_shape, + block_idx=compute_block_idx, + smem_workspace=self._smem_workspace, + smem_base=smem_base, + device_workspace=self._device_workspace, + is_fallback_cluster=is_fallback_cluster, + ) + scheduler = self.scheduler + + fc2_spin_threshold = (intermediate_gateup + self._mainloop.cta_tile_m - 1) // self._mainloop.cta_tile_m + helper_weight_ready_gate = None + if cutlass.const_expr(self.helper_expert_count > 0): + helper_weight_ready_gate = make_helper_weight_ready_generation_gate( + memory_slot_count=self.local_expert_count, + helper_count=self.helper_expert_count, + source_count=self.world_size, + terminal_flags=hot_expert_weight_ready_flags, + expected_generation=hot_expert_weight_ready_generation, + ) + kernel_extension = BlockScaledSwapAbFc12Extension( + sf_vec_size=self.sf_vec_size, + fc1_done_counter_pointer=fc1_done_counter.iterator, + fc2_spin_threshold=fc2_spin_threshold, + fc1_ready_counter_pointer=self.token_comm.fc1_ready_counter_pointer(self._device_workspace), + ) + optional_epilogue_args = GatedActEpilogueArgs( + fc1_alpha=fc1_alpha, fc2_alpha=fc2_alpha, fc1_norm_const=fc1_norm_const, topk_scores=pool_topk_scores + ) + + pipeline_init_arrive_mixed_cga( + self.cluster_shape_mn, self.resolved_fallback_cluster_shape_mn, is_fallback_cluster, is_relaxed=True + ) + pipeline_init_wait_mixed_cga( + self.cluster_shape_mn, self.resolved_fallback_cluster_shape_mn, is_fallback_cluster + ) + + if warp_idx == self.scheduler_warp_id: + cute.arch.warpgroup_reg_dealloc(self.scheduler_register_count) + iket.range_push("mega.scheduler") + if cutlass.const_expr(hasattr(scheduler, "initialize_fallback_group")): + iket.range_push("scheduler.register_group") + scheduler.initialize_fallback_group() + iket.range_pop() + iket.range_push("scheduler.wait_sizes_ready") + self.token_comm.wait_for_sizes_ready(self._device_workspace) + iket.range_pop() + iket.range_push("scheduler.gen_work") + work_tile = scheduler.gen_next_work() + iket.range_pop() + helper_ready_latched = Int32(0) + while work_tile.is_valid_tile: + iket.range_push("scheduler.publish_work") + prepared_tile = kernel_extension.prepare_work_tile(work_tile) + if cutlass.const_expr(helper_weight_ready_gate is not None): + helper_ready_latched = helper_weight_ready_gate.wait_before_first_helper_publish( + prepared_tile.expert_idx, helper_ready_latched + ) + scheduler.publish_work(prepared_tile) + iket.range_pop() + iket.range_push("scheduler.gen_work") + work_tile = scheduler.gen_next_work() + iket.range_pop() + iket.range_push("scheduler.publish_tail") + scheduler.publish_work(work_tile) + scheduler.produce_tail() + iket.range_pop() + iket.range_pop() + + if warp_idx == self.tma_a_warp_id: + cute.arch.warpgroup_reg_dealloc(self.other_warp_register_count) + iket.range_push("mega.tma_a") + sched_consumer = scheduler.make_consumer() + self._mainloop.run_tma_a( + fc1_tma_a_tensor=fc1_tma_a_tensor, + fc1_tma_a_atom=fc1_tma_a_atom, + fc1_tma_sfa_tensor=fc1_tma_sfa_tensor, + fc1_tma_sfa_atom=fc1_tma_sfa_atom, + fc2_tma_a_tensor=fc2_tma_a_tensor, + fc2_tma_a_atom=fc2_tma_a_atom, + fc2_tma_sfa_tensor=fc2_tma_sfa_tensor, + fc2_tma_sfa_atom=fc2_tma_sfa_atom, + ab_pipeline=ab_pipeline, + ab_pipeline_state=tma_a_pipeline_state, + sched_consumer=sched_consumer, + kernel_extension=kernel_extension, + ) + iket.range_pop() + + if warp_idx == self.tma_b_warp_id: + cute.arch.warpgroup_reg_dealloc(self.other_warp_register_count) + iket.range_push("mega.tma_b") + sched_consumer = scheduler.make_consumer() + self._mainloop.run_tma_b( + fc1_tma_b_tensor=fc1_tma_b_tensor, + fc1_tma_b_atom_or_pair=fc1_tma_b_atom_or_pair, + fc1_tma_sfb_tensor=fc1_tma_sfb_tensor, + fc1_tma_sfb_atom_or_pair=fc1_tma_sfb_atom_or_pair, + fc2_tma_b_tensor=fc2_tma_b_tensor, + fc2_tma_b_atom_or_pair=fc2_tma_b_atom_or_pair, + fc2_tma_sfb_tensor=fc2_tma_sfb_tensor, + fc2_tma_sfb_atom_or_pair=fc2_tma_sfb_atom_or_pair, + ab_pipeline=ab_pipeline, + ab_pipeline_state=tma_b_pipeline_state, + sched_consumer=sched_consumer, + kernel_extension=kernel_extension, + fc1_done_counter_pointer=fc1_done_counter.iterator, + fc2_spin_threshold=fc2_spin_threshold, + ) + iket.range_pop() + + if warp_idx == self.mma_warp_id: + cute.arch.warpgroup_reg_dealloc(self.other_warp_register_count) + iket.range_push("mega.mma") + sched_consumer = scheduler.make_consumer() + self._mainloop.run_mma( + tmem_allocator=tmem_allocator, + ab_pipeline=ab_pipeline, + ab_pipeline_state=mma_pipeline_state, + acc_pipeline=acc_pipeline, + sched_consumer=sched_consumer, + ) + iket.range_pop() + + if warp_idx < len(self.epilogue_warp_ids): + cute.arch.warpgroup_reg_alloc(self.epilogue_register_count) + iket.range_push("mega.epilogue") + sched_consumer = scheduler.make_consumer() + tmem_allocator.allocate(self._mainloop.num_tmem_alloc_cols) + tmem_allocator.wait_for_alloc() + tmem_pointer = tmem_allocator.retrieve_ptr(self.acc_dtype) + self.epilogue.run( + self._smem_workspace, + smem_base, + tmem_pointer, + acc_pipeline, + sched_consumer, + kernel_extension, + fc1_output_tma_atom, + fc1_output_tma_tensor, + fc1_output_sf, + fc2_output_tma_atom, + fc2_output_tma_tensor, + fc2_output, + fc1_done_counter, + thread_idx, + token_src_metadata, + fc2_done_counter, + fc2_output_sf, + peer_rank_ptr_mapper, + optional_epilogue_args, + ) + tmem_allocator.relinquish_alloc_permit() + tmem_allocator.free(tmem_allocator.retrieve_ptr(self.acc_dtype), self._mainloop.num_tmem_alloc_cols) + iket.range_pop() + + # Donate registers from complete warpgroups, including idle warps. + if warp_idx >= Int32(self.transfer_warp_idx_start): + cute.arch.warpgroup_reg_dealloc(self.other_warp_register_count) + if cutlass.const_expr(self.token_back_on_compute_ctas): + if (warp_idx >= self.token_comm.combine_warp_start) & ( + warp_idx < self.token_comm.combine_warp_start + self.token_comm.combine_warp_count + ): + self.token_comm.token_back(self._smem_workspace, smem_base) + + # Every physical CTA participates in the common tail, including idle transfer warps. + # Each operation has already drained its own asynchronous copies before this CTA barrier. + iket.range_push("mega.kernel_tail_wait") + cute.arch.sync_threads() + iket.range_pop() + if (warp_idx >= Int32(self.transfer_warp_idx_start)) & (warp_idx < Int32(self.token_in_end_warp_idx)): + iket.range_push("mega.tail_reset") + self.token_comm.reset_tail() + iket.range_pop() + + self.token_comm.remove_device_members() + workspace.remove_device_members() diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py index 2c864e63cc38..c5567f3df5a0 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py @@ -1,7 +1,5 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: BSD-3-Clause -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: BSD-3-Clause """Generation-phase MegaMoE kernel composition for Rubin. Specialised for few tokens per rank and large expert parallelism. The narrowing @@ -27,10 +25,12 @@ workspace or exits before the clearing is visible. """ +import hashlib import math from typing import ClassVar, Optional, Tuple import cuda.bindings.driver as cuda + import cutlass import cutlass.cute as cute import cutlass.pipeline as pipeline @@ -40,36 +40,28 @@ from cutlass.cutlass_dsl import Int32, Int64 from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait -from .....api import ( - ImplDesc, - KernelClass, - OptionalRequirement, - ProblemDesc, - StaticOrRuntimeIntegerType, -) +from .....api import ImplDesc, KernelClass, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType from .....communication.nvlink_domain.symmetric_buffer import SymmetricBufferDevice -from .....helpers.cute_py_helpers import ( - Tcgen05MmaInstruction, - make_tcgen05_tmem_plan, - tcgen05_block_scaled_acc_dtype, -) +from .....helpers.helper_weight_ready_gate import derive_helper_weight_ready_topology, make_helper_weight_ready_generation_gate +from .....helpers.cute_py_helpers import Tcgen05MmaInstruction, make_tcgen05_tmem_plan, tcgen05_block_scaled_acc_dtype from .....helpers.device_workspace import DeviceWorkspace +from .....helpers.dlb_expert_count import resolve_dlb_expert_counts from .....helpers.dsl_helpers import spin_wait from .....helpers.iket_compat import iket +from .....helpers.megamoe_aot import make_megamoe_aot_callable from .....helpers.smem_workspace import SmemWorkspace from .....helpers.software_sync import NvlinkBarrier +from .....helpers.topk_reduce import TopkReduce from .....helpers.utils import ceil_div, round_up from .....quant_def import CombineFormat, QuantKind -from ....blackwell.inference.mega.topk_reduce import TopkReduce from ....schedulers.base import WorkIdAcquisitionMode from ....schedulers.fc12_scheduler import BlackwellFusedFc12Scheduler, PhaseInterleavedFc12Scheduler from .block_scaled_swap_ab_fc12_epilogue import GatedActEpilogueArgs, SwapABGatedActEpilogue from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension -from .block_scaled_swap_ab_fc12_mainloop_gen_specialized import ( - BlockScaledSwapAbFc12MainloopGenSpecialized, -) +from .block_scaled_swap_ab_fc12_mainloop_gen_specialized import BlockScaledSwapAbFc12MainloopGenSpecialized from .token_comm_gen_specialized import GenphaseTokenComm, GenphaseTokenCommArgs + _aot_symbol_prefix = "rubin_genphase_mega_moe_aot" @@ -140,10 +132,15 @@ def impl_desc_require(cls) -> dict[str, object]: "epi_flag_batches": tuple, "launch_cluster_count": int, "reduce_topk_in_kernel": bool, + "helper_expert_count": OptionalRequirement(int), + "total_helper_slots": OptionalRequirement(int), } def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self._validate_desc_inputs(problem_desc, impl_desc) + problem_desc, impl_desc = resolve_dlb_expert_counts( + problem_desc, impl_desc, world_size=problem_desc["world_size"] + ) for field_name, fixed_value, supplied_value in ( ("cluster_shape_mn", self.cluster_shape_mn, impl_desc.get("cluster_shape_mn")), @@ -157,10 +154,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: ) fallback_cluster_shape_mn = impl_desc.get("fallback_cluster_shape_mn") - if ( - fallback_cluster_shape_mn is not None - and fallback_cluster_shape_mn != self.cluster_shape_mn - ): + if fallback_cluster_shape_mn is not None and fallback_cluster_shape_mn != self.cluster_shape_mn: raise ValueError( f"The generation-phase kernel launches a single {self.cluster_shape_mn} cluster shape, so " f"fallback_cluster_shape_mn must be unset or equal to it, got {fallback_cluster_shape_mn}." @@ -199,11 +193,19 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.epi_flag_batches = impl_desc["epi_flag_batches"] self.launch_cluster_count = impl_desc["launch_cluster_count"] self.reduce_topk_in_kernel = impl_desc["reduce_topk_in_kernel"] - + self.helper_expert_count = impl_desc.get("helper_expert_count", 0) + if type(self.helper_expert_count) is not int or self.helper_expert_count < 0: + raise ValueError("helper_expert_count must be a non-negative exact int") self.occupancy = 1 self.architecture = "sm_107" self.threads_per_cta = 12 * 32 self.local_expert_count = self.expert_count // self.world_size + if self.helper_expert_count > 0: + self.home_expert_count = derive_helper_weight_ready_topology( + self.local_expert_count, self.helper_expert_count + ).home_expert_count + else: + self.home_expert_count = self.local_expert_count self.cluster_size = self.cluster_shape_mn[0] * self.cluster_shape_mn[1] self.promised_launchable_sm_count = self.launch_cluster_count * self.cluster_size @@ -235,9 +237,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.token_comm = GenphaseTokenComm(problem_desc, token_comm_impl_desc) self.max_routed_rows = self.token_comm.worst_case_data_rows self._validate_communication_geometry() - self.tail_barrier = NvlinkBarrier( - world_size=self.world_size, barrier_id=self.tail_barrier_id - ) + self.tail_barrier = NvlinkBarrier(world_size=self.world_size, barrier_id=self.tail_barrier_id) fc12_problem_desc = ProblemDesc({**problem_desc, "expert_count": self.local_expert_count}) resolved_impl_desc = ImplDesc( @@ -263,9 +263,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: } ) scheduler_type = ( - BlackwellFusedFc12Scheduler - if self.schedule_mode == "grouped" - else PhaseInterleavedFc12Scheduler + BlackwellFusedFc12Scheduler if self.schedule_mode == "grouped" else PhaseInterleavedFc12Scheduler ) self.scheduler = scheduler_type(fc12_problem_desc, resolved_impl_desc) self.epilogue = SwapABGatedActEpilogue(fc12_problem_desc, resolved_impl_desc) @@ -281,17 +279,11 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: "selects the peer-facing UBLK store." ) if self.epilogue.token_back_enabled: - raise ValueError( - "The generation-phase communication component provides no token-back path." - ) + raise ValueError("The generation-phase communication component provides no token-back path.") self._device_workspace = self._build_device_workspace() - self._mainloop, self._smem_workspace = self._build_mainloop_and_smem( - fc12_problem_desc, resolved_impl_desc - ) + self._mainloop, self._smem_workspace = self._build_mainloop_and_smem(fc12_problem_desc, resolved_impl_desc) self._topk_reduce = ( - None - if self.reduce_topk_in_kernel - else TopkReduce(self.hidden_size, self.topk, self.combine_format) + None if self.reduce_topk_in_kernel else TopkReduce(self.hidden_size, self.topk, self.combine_format) ) # ------------------------------------------------------------------ @@ -305,13 +297,9 @@ def _validate_geometry(self) -> None: isinstance(self.hidden_size, int), ) if any(static_expert_dimensions) and not all(static_expert_dimensions): - raise ValueError( - "Genphase MegaMoE expert dimensions must be either all static or all runtime." - ) + raise ValueError("Genphase MegaMoE expert dimensions must be either all static or all runtime.") if not all(static_expert_dimensions): - raise NotImplementedError( - "The Genphase MegaMoE kernel currently requires static expert dimensions." - ) + raise NotImplementedError("The Genphase MegaMoE kernel currently requires static expert dimensions.") if self.expert_count <= 0 or self.intermediate_gateup_size <= 0 or self.hidden_size <= 0: raise ValueError("Genphase MegaMoE expert dimensions must be positive.") if self.world_size <= 0 or self.expert_count % self.world_size != 0: @@ -319,9 +307,7 @@ def _validate_geometry(self) -> None: if self.topk <= 0 or self.topk > self.expert_count: raise ValueError("topk must be positive and no greater than expert_count.") if self.topk_index_dtype not in (cutlass.Int32, cutlass.Int64): - raise ValueError( - f"topk_index_dtype must be Int32 or Int64, got {self.topk_index_dtype}." - ) + raise ValueError(f"topk_index_dtype must be Int32 or Int64, got {self.topk_index_dtype}.") # The FC2 epilogue carries no per-row top-k weight, so folding the weight # in at FC1 is the only way an in-kernel reduction can be correct. The # alternative costs a (token, topk, hidden) symmetric plane and a second @@ -352,20 +338,14 @@ def _validate_geometry(self) -> None: raise ValueError( f"schedule_policy mode must be 'grouped' or 'phase_interleave', got {self.schedule_mode!r}." ) - if self.hint is not None and ( - isinstance(self.hint, bool) or not isinstance(self.hint, int) or self.hint <= 0 - ): - raise ValueError( - f"schedule_policy hint must be a positive Python int, got {self.hint!r}." - ) + if self.hint is not None and (isinstance(self.hint, bool) or not isinstance(self.hint, int) or self.hint <= 0): + raise ValueError(f"schedule_policy hint must be a positive Python int, got {self.hint!r}.") if self.schedule_mode == "phase_interleave" and self.work_id_mode != "atomic_counter": raise ValueError("phase_interleave currently requires work_id_mode='atomic_counter'.") if len(self.epi_flag_batches) != 2: raise ValueError("epi_flag_batches must contain FC1 and FC2 batch sizes.") if any(batch < 1 or batch > len(self.epilogue_warp_ids) for batch in self.epi_flag_batches): - raise ValueError( - f"epi_flag_batches values must be in [1, {len(self.epilogue_warp_ids)}]." - ) + raise ValueError(f"epi_flag_batches values must be in [1, {len(self.epilogue_warp_ids)}].") if self.sf_vec_size <= 0: raise ValueError("sf_vec_size must be positive.") @@ -375,10 +355,7 @@ def _validate_geometry(self) -> None: ): if len(dimensions) != expected_rank: raise ValueError(f"{field_name} must contain {expected_rank} dimensions.") - if not all( - isinstance(dimension, int) and not isinstance(dimension, bool) - for dimension in dimensions - ): + if not all(isinstance(dimension, int) and not isinstance(dimension, bool) for dimension in dimensions): raise TypeError(f"{field_name} dimensions must be Python integers.") if any(dimension <= 0 for dimension in dimensions): raise ValueError(f"{field_name} dimensions must be positive.") @@ -398,9 +375,7 @@ def _validate_geometry(self) -> None: f"{self.quant_kind} requires Rubin 2x instruction K={expected_instruction_k}, got {instruction_k}." ) if (instruction_m, instruction_n) != (mma_m, mma_n): - raise NotImplementedError( - "Rubin Genphase MegaMoE does not implement M/N instruction repetition." - ) + raise NotImplementedError("Rubin Genphase MegaMoE does not implement M/N instruction repetition.") # Tile N 64 is excluded because it turns on the halved SFB tile index and # the SFB TMEM shift, and 512 is excluded because the FC1 gather needs one # B row per 128-byte swizzle atom (see the mainloop's gather geometry). @@ -411,13 +386,9 @@ def _validate_geometry(self) -> None: if mma_k % (self.sf_vec_size * 4) != 0: raise ValueError("mma_tiler K must be divisible by four scale-factor vectors.") if self.intermediate_gateup_size % (self.sf_vec_size * 4) != 0: - raise ValueError( - "The intermediate dimension must be divisible by four scale-factor vectors." - ) + raise ValueError("The intermediate dimension must be divisible by four scale-factor vectors.") if mma_m // mma_cta_count != 128: - raise ValueError( - f"The epilogue fixes the per-CTA output tile at 128, got {mma_m // mma_cta_count}." - ) + raise ValueError(f"The epilogue fixes the per-CTA output tile at 128, got {mma_m // mma_cta_count}.") # The epilogue picks the TMA owner warp with `warp_idx == subtile_idx` and # does not check the bound itself, so a wider token tile would drop stores. subtile_count = mma_n // 64 @@ -453,9 +424,7 @@ def _resolve_schedule_hint(self) -> None: blocks_fc1 = ceil_div(self.intermediate_gateup_size, cluster_feature_tile) blocks_fc2 = ceil_div(self.hidden_size, cluster_feature_tile) # Cover the worst token-block alignment of one full persistent-cluster FC2 claim wave. - max_dependent_token_blocks = ceil_div( - self.launch_cluster_count + blocks_fc2 - 1, blocks_fc2 - ) + max_dependent_token_blocks = ceil_div(self.launch_cluster_count + blocks_fc2 - 1, blocks_fc2) required_fc1_work = max_dependent_token_blocks * blocks_fc1 raw_minimum_hint = max(1, ceil_div(required_fc1_work, self.launch_cluster_count)) minimum_safe_hint = round_up(raw_minimum_hint, self.epi_flag_batches[0]) @@ -477,9 +446,7 @@ def _build_device_workspace(self) -> DeviceWorkspace: max_sf_rows = self.token_comm.worst_case_sf_rows # One slot per token block across every local expert, which is the space # `cumulative_token_block_count + tile_n_idx` indexes. - counter_slot_count = ( - self.token_comm.worst_case_padded_rows(self.cta_tile_n) // self.cta_tile_n - ) + counter_slot_count = self.token_comm.worst_case_padded_rows(self.cta_tile_n) // self.cta_tile_n device_workspace = DeviceWorkspace() device_workspace.register( @@ -529,7 +496,7 @@ def _build_mainloop_and_smem( ) -> Tuple[BlockScaledSwapAbFc12MainloopGenSpecialized, SmemWorkspace]: # The communication component builds its own SMEM workspace for its own # launch, so nothing of it is charged against this kernel's budget. - smem_limit = cutlass.memory.get_smem_capacity_in_bytes(self.architecture) // self.occupancy + smem_limit = utils.get_smem_capacity_in_bytes(self.architecture) // self.occupancy smem_workspace = SmemWorkspace() self.scheduler.register_smem_regions(smem_workspace) self.epilogue.register_smem_regions(smem_workspace) @@ -539,9 +506,7 @@ def _build_mainloop_and_smem( if mainloop_smem_budget_bytes <= 0: raise ValueError("Fixed kernel components leave no SMEM budget for the Mainloop.") - mainloop_impl_desc = ImplDesc( - {**resolved_impl_desc, "mainloop_smem_budget_bytes": mainloop_smem_budget_bytes} - ) + mainloop_impl_desc = ImplDesc({**resolved_impl_desc, "mainloop_smem_budget_bytes": mainloop_smem_budget_bytes}) mainloop = BlockScaledSwapAbFc12MainloopGenSpecialized(problem_desc, mainloop_impl_desc) mainloop.register_smem_regions(smem_workspace) @@ -573,7 +538,7 @@ def dtype_name(dtype: type) -> str: f"sm107_block_scaled_swap_ab_gen_phase_specialized_moe_{self.quant_kind}_" f"{dtype_name(self.a_dtype)}_{dtype_name(self.b_dtype)}_{dtype_name(self.acc_dtype)}_" f"sfvec{self.sf_vec_size}_a{self.a_major_mode.name.lower()}_b{self.b_major_mode.name.lower()}_" - f"w{self.world_size}_e{self.expert_count}_topk{self.topk}_" + f"w{self.world_size}_e{self.expert_count}_topk{self.topk}_helpers{self.helper_expert_count}_" f"topkidx{dtype_name(self.topk_index_dtype)}_" f"h{self.hidden_size}_i{self.intermediate_gateup_size}_maxtoken{self.max_tokens_per_rank}_" f"inst{instruction}_{self.mma_k_mode}_tile{tile}_" @@ -589,6 +554,7 @@ def dtype_name(dtype: type) -> str: f"{'apply_topk_fc1' if self.apply_topk_at_fc1 else 'apply_topk_fc2'}_" f"{'inkernel_reduce' if self.reduce_topk_in_kernel else 'separate_reduce'}_" f"abstages{self._mainloop.num_weight_ab_stages}x{self._mainloop.num_token_ab_stages}" + + ("_dlb_v1" if self.helper_expert_count > 0 else "") ) def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): @@ -606,14 +572,10 @@ def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): extents = tuple( - sym_int64(divisibility=math.gcd(int(extent), 128)) - if axis in dynamic_axes - else int(extent) + sym_int64(divisibility=math.gcd(int(extent), 128)) if axis in dynamic_axes else int(extent) for axis, extent in enumerate(shape) ) - return make_fake_compact_tensor( - dtype, extents, stride_order=stride_order, assumed_align=alignment - ) + return make_fake_compact_tensor(dtype, extents, stride_order=stride_order, assumed_align=alignment) tokens = self.max_tokens_per_rank hidden = self.hidden_size @@ -622,12 +584,8 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): experts = self.local_expert_count sf_vec_size = self.sf_vec_size # Weight SF is atom-swizzled and opaque, so its extent is whatever the caller staged. - fc1_weight_sf_columns = round_up(intermediate_gateup, 128) * round_up( - hidden // sf_vec_size, 4 - ) - fc2_weight_sf_columns = round_up(hidden, 128) * round_up( - intermediate_downproj // sf_vec_size, 4 - ) + fc1_weight_sf_columns = round_up(intermediate_gateup, 128) * round_up(hidden // sf_vec_size, 4) + fc2_weight_sf_columns = round_up(hidden, 128) * round_up(intermediate_downproj // sf_vec_size, 4) fake_arguments = dict( activation=fake_tensor( @@ -646,18 +604,10 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): ), topk_indices=fake_tensor(self.topk_index_dtype, (tokens, self.topk), (1, 0), {}, 16), topk_scores=fake_tensor(cutlass.Float32, (tokens, self.topk), (1, 0), {}, 4), - fc1_weight=fake_tensor( - self.a_dtype, (experts, hidden, intermediate_gateup), (2, 0, 1), {1, 2}, 16 - ), - fc1_weight_sf=fake_tensor( - self.sf_dtype, (experts, fc1_weight_sf_columns), (1, 0), {1}, 16 - ), - fc2_weight=fake_tensor( - self.a_dtype, (experts, intermediate_downproj, hidden), (2, 0, 1), {1, 2}, 16 - ), - fc2_weight_sf=fake_tensor( - self.sf_dtype, (experts, fc2_weight_sf_columns), (1, 0), {1}, 16 - ), + fc1_weight=fake_tensor(self.a_dtype, (experts, hidden, intermediate_gateup), (2, 0, 1), {1, 2}, 16), + fc1_weight_sf=fake_tensor(self.sf_dtype, (experts, fc1_weight_sf_columns), (1, 0), {1}, 16), + fc2_weight=fake_tensor(self.a_dtype, (experts, intermediate_downproj, hidden), (2, 0, 1), {1, 2}, 16), + fc2_weight_sf=fake_tensor(self.sf_dtype, (experts, fc2_weight_sf_columns), (1, 0), {1}, 16), output_activation=fake_tensor(cutlass.BFloat16, (tokens, hidden), (1, 0), {}, 16), local_workspace=make_ptr(cutlass.Uint8, 0, AddressSpace.gmem, assumed_align=128), shared_workspace=make_ptr(cutlass.Uint8, 0, AddressSpace.gmem, assumed_align=128), @@ -672,6 +622,11 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): ), stream=make_fake_stream(), ) + if self.helper_expert_count > 0: + fake_arguments.update( + hot_expert_weight_ready_flags=make_ptr(cutlass.Uint64, 0, AddressSpace.gmem, assumed_align=8), + hot_expert_weight_ready_generation=cutlass.Uint64(1), + ) if self.quant_kind.uses_global_scale: fake_arguments.update( fc1_alpha=fake_tensor(cutlass.Float32, (experts,), (0,), set(), 4), @@ -682,16 +637,18 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): compiled = cute.compile[cute.EnableTVMFFI(True)](self, **fake_arguments) if out_path is None: return compiled - compiled.export_to_c( - out_path, function_name=_aot_symbol_prefix, export_only_tvm_ffi_symbols=True - ) + compiled.export_to_c(out_path, function_name=self._aot_symbol_name(), export_only_tvm_ffi_symbols=True) return out_path - @staticmethod - def load_compiled(path: str): + def _aot_symbol_name(self) -> str: + identity = hashlib.sha256(self.name().encode()).hexdigest()[:16] + return f"{_aot_symbol_prefix}_{identity}" + + def load_compiled(self, path: str): from cutlass.cute.runtime import load_module - return load_module(path, enable_tvm_ffi=True)[_aot_symbol_prefix] + function = getattr(load_module(path, enable_tvm_ffi=True), self._aot_symbol_name()) + return make_megamoe_aot_callable(function, self) # ------------------------------------------------------------------ # Host entry @@ -716,16 +673,23 @@ def __call__( fc1_alpha: Optional[cute.Tensor] = None, # (local_experts,) fc2_alpha: Optional[cute.Tensor] = None, # (local_experts,) fc1_norm_const: Optional[cute.Tensor] = None, # (local_experts,) + hot_expert_weight_ready_flags: Optional[cute.Pointer] = None, # uint64[EP] + hot_expert_weight_ready_generation: Optional[cutlass.Uint64] = None, ) -> None: """Launch generation-phase dispatch, fused FC12, and the optional top-k reduce.""" + helper_ready_bundle = (hot_expert_weight_ready_flags, hot_expert_weight_ready_generation) + if cutlass.const_expr( + (self.helper_expert_count > 0 and not all(value is not None for value in helper_ready_bundle)) + or (self.helper_expert_count == 0 and any(value is not None for value in helper_ready_bundle)) + ): + raise ValueError("helper READY flags/generation must both be present iff helper_expert_count > 0") + def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: return cute.make_tensor(tensor.iterator, cute.make_layout(shape, stride=tensor.stride)) if cutlass.const_expr(topk_indices.element_type is not self.topk_index_dtype): - raise TypeError( - f"topk_indices must be the declared {self.topk_index_dtype}, got {topk_indices.dtype}." - ) + raise TypeError(f"topk_indices must be the declared {self.topk_index_dtype}, got {topk_indices.dtype}.") for operand_name, operand, expected_dtype in ( ("activation", activation, self.b_dtype), ("activation_sf", activation_sf, self.sf_dtype), @@ -767,12 +731,8 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: self._device_workspace.assign_device_members(local_workspace, shared_workspace) peer_rank_ptr_mapper = peer_rank_ptr_mapper_host.make_device_object() dense_activation = self.token_comm.dense_activation_tensor(self._device_workspace) - canonical_activation_sf = self.token_comm.canonical_activation_sf_tensor( - self._device_workspace - ) - scheduler_expert_sizes = self.token_comm.scheduler_expert_sizes_tensor( - self._device_workspace - ) + canonical_activation_sf = self.token_comm.canonical_activation_sf_tensor(self._device_workspace) + scheduler_expert_sizes = self.token_comm.scheduler_expert_sizes_tensor(self._device_workspace) fc1_output = self._device_workspace.tensor(self.fc1_output_region) fc1_output_sf = self._device_workspace.tensor(self.fc1_output_sf_region) @@ -821,20 +781,14 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: padded_intermediate_gateup = cute.round_up(intermediate_gateup, sf_vec_size * 4) expected_fc1_weight_sf_columns = padded_intermediate_gateup * hidden // sf_vec_size if cutlass.const_expr( - isinstance(fc1_weight_sf.shape[1], int) - and isinstance(expected_fc1_weight_sf_columns, int) + isinstance(fc1_weight_sf.shape[1], int) and isinstance(expected_fc1_weight_sf_columns, int) ): if cutlass.const_expr(fc1_weight_sf.shape[1] != expected_fc1_weight_sf_columns): raise ValueError("fc1_weight_sf has an incompatible column count.") fc1_sfa = cute.make_tensor( fc1_weight_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - cutlass.Int32(padded_intermediate_gateup), - cutlass.Int32(hidden), - cutlass.Int32(experts), - ), - sf_vec_size, + (cutlass.Int32(padded_intermediate_gateup), cutlass.Int32(hidden), cutlass.Int32(experts)), sf_vec_size ), ) @@ -851,31 +805,21 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: padded_fc2_intermediate = cute.round_up(intermediate_fc2, sf_vec_size * 4) expected_fc2_weight_sf_columns = padded_fc2_hidden * padded_fc2_intermediate // sf_vec_size if cutlass.const_expr( - isinstance(fc2_weight_sf.shape[1], int) - and isinstance(expected_fc2_weight_sf_columns, int) + isinstance(fc2_weight_sf.shape[1], int) and isinstance(expected_fc2_weight_sf_columns, int) ): if cutlass.const_expr(fc2_weight_sf.shape[1] != expected_fc2_weight_sf_columns): raise ValueError("fc2_weight_sf has an incompatible column count.") fc2_sfa = cute.make_tensor( fc2_weight_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - cutlass.Int32(padded_fc2_hidden), - cutlass.Int32(padded_fc2_intermediate), - cutlass.Int32(experts_fc2), - ), + (cutlass.Int32(padded_fc2_hidden), cutlass.Int32(padded_fc2_intermediate), cutlass.Int32(experts_fc2)), sf_vec_size, ), ) fc2_sfb = cute.make_tensor( fc1_output_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - fc1_output_sf.shape[0], - cutlass.Int32(fc1_output_sf.shape[1] * sf_vec_size), - singleton, - ), - sf_vec_size, + (fc1_output_sf.shape[0], cutlass.Int32(fc1_output_sf.shape[1] * sf_vec_size), singleton), sf_vec_size ), ) @@ -887,20 +831,13 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: output_activation.iterator, cute.make_layout( (self.max_tokens_per_rank, 1, hidden), - stride=( - output_activation.stride[0], - output_activation.stride[0], - output_activation.stride[1], - ), + stride=(output_activation.stride[0], output_activation.stride[0], output_activation.stride[1]), ), ) else: - pre_reduced_activation = self._device_workspace.tensor( - self.pre_reduced_activation_region - ) + pre_reduced_activation = self._device_workspace.tensor(self.pre_reduced_activation_region) fc2_output = rewrite_tensor_shape( - pre_reduced_activation, - (pre_reduced_activation.shape[0], pre_reduced_activation.shape[1], hidden), + pre_reduced_activation, (pre_reduced_activation.shape[0], pre_reduced_activation.shape[1], hidden) ) self._mainloop.materialize_codegen_members() @@ -966,6 +903,8 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: fc1_alpha, fc2_alpha, fc1_norm_const, + hot_expert_weight_ready_flags, + hot_expert_weight_ready_generation, ).launch( grid=grid, block=[self.threads_per_cta, 1, 1], @@ -976,9 +915,7 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: ) if cutlass.const_expr(not self.reduce_topk_in_kernel): reduce_scores = None if cutlass.const_expr(self.apply_topk_at_fc1) else topk_scores - self._topk_reduce( - pre_reduced_activation, None, output_activation, reduce_scores, stream - ) + self._topk_reduce(pre_reduced_activation, None, output_activation, reduce_scores, stream) # ------------------------------------------------------------------ # Device entry @@ -1015,6 +952,8 @@ def _kernel( fc1_alpha: Optional[cute.Tensor], fc2_alpha: Optional[cute.Tensor], fc1_norm_const: Optional[cute.Tensor], + hot_expert_weight_ready_flags: Optional[cute.Pointer], + hot_expert_weight_ready_generation: Optional[cutlass.Uint64], ): """Compose scheduler, mainloop, epilogue, and the workspace tail reset.""" self._mainloop.materialize_codegen_members() @@ -1037,8 +976,7 @@ def _kernel( fc1_output_sf = cute.make_tensor( fc1_output_sf_storage.iterator, blockscaled_utils.tile_atom_to_shape_SF( - (fc1_output_sf_storage.shape[0], intermediate_gateup // 2, cutlass.Int32(1)), - self.epilogue.sf_vec_size, + (fc1_output_sf_storage.shape[0], intermediate_gateup // 2, cutlass.Int32(1)), self.epilogue.sf_vec_size ), ) @@ -1097,9 +1035,7 @@ def _kernel( ) scheduler = self.scheduler - fc2_spin_threshold = ( - intermediate_gateup + self._mainloop.cta_tile_m - 1 - ) // self._mainloop.cta_tile_m + fc2_spin_threshold = (intermediate_gateup + self._mainloop.cta_tile_m - 1) // self._mainloop.cta_tile_m # The scheduler enumerates slots, so the tile it produces has to be # rewritten into an expert index before any consumer sees it: the weights, # their scale factors, the per-expert alphas and the per-expert @@ -1108,6 +1044,15 @@ def _kernel( expert_remap = self.token_comm.slot_to_expert_tensor(self._device_workspace) # No FC1 ready counter: readiness is published per rank, not per tile, and # the waits below cover it once per warp. + helper_weight_ready_gate = None + if cutlass.const_expr(self.helper_expert_count > 0): + helper_weight_ready_gate = make_helper_weight_ready_generation_gate( + memory_slot_count=self.local_expert_count, + helper_count=self.helper_expert_count, + source_count=self.world_size, + terminal_flags=hot_expert_weight_ready_flags, + expected_generation=hot_expert_weight_ready_generation, + ) kernel_extension = BlockScaledSwapAbFc12Extension( sf_vec_size=self.sf_vec_size, fc1_done_counter_pointer=fc1_done_counter.iterator, @@ -1116,10 +1061,7 @@ def _kernel( expert_remap=expert_remap, ) optional_epilogue_args = GatedActEpilogueArgs( - fc1_alpha=fc1_alpha, - fc2_alpha=fc2_alpha, - fc1_norm_const=fc1_norm_const, - topk_scores=pool_topk_scores, + fc1_alpha=fc1_alpha, fc2_alpha=fc2_alpha, fc1_norm_const=fc1_norm_const, topk_scores=pool_topk_scores ) pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) @@ -1143,9 +1085,15 @@ def _kernel( iket.range_push("scheduler.gen_work") work_tile = scheduler.gen_next_work() iket.range_pop() + helper_ready_latched = Int32(0) while work_tile.is_valid_tile: iket.range_push("scheduler.publish_work") - scheduler.publish_work(kernel_extension.prepare_work_tile(work_tile)) + prepared_tile = kernel_extension.prepare_work_tile(work_tile) + if cutlass.const_expr(helper_weight_ready_gate is not None): + helper_ready_latched = helper_weight_ready_gate.wait_before_first_helper_publish( + prepared_tile.expert_idx, helper_ready_latched + ) + scheduler.publish_work(prepared_tile) iket.range_pop() iket.range_push("scheduler.gen_work") work_tile = scheduler.gen_next_work() @@ -1193,9 +1141,7 @@ def _kernel( kernel_extension=kernel_extension, fc1_done_counter_pointer=fc1_done_counter.iterator, fc2_spin_threshold=fc2_spin_threshold, - fc1_sf_ready_pointer=self._device_workspace.ptr( - self.token_comm.fc1_sf_ready_region - ), + fc1_sf_ready_pointer=self._device_workspace.ptr(self.token_comm.fc1_sf_ready_region), expert_sizes=self._device_workspace.tensor(self.token_comm.expert_sizes_region), ) iket.range_pop() @@ -1262,30 +1208,21 @@ def _kernel( optional_epilogue_args, ) tmem_allocator.relinquish_alloc_permit() - tmem_allocator.free( - tmem_allocator.retrieve_ptr(self.acc_dtype), self._mainloop.num_tmem_alloc_cols - ) + tmem_allocator.free(tmem_allocator.retrieve_ptr(self.acc_dtype), self._mainloop.num_tmem_alloc_cols) epilogue_thread_count = 32 * len(self.epilogue_warp_ids) iket.range_push("tail.grid_sync_before_reset") self.tail_barrier.sync( - epilogue_thread_count, - Int32(self.promised_launchable_sm_count), - linear_cta_idx, - thread_idx, + epilogue_thread_count, Int32(self.promised_launchable_sm_count), linear_cta_idx, thread_idx ) iket.range_pop() iket.range_push("tail.reset_workspace") total_reset_threads = self.promised_launchable_sm_count * epilogue_thread_count global_reset_thread = linear_cta_idx * Int32(epilogue_thread_count) + thread_idx - self._device_workspace.reset_tail_space( - "shared", global_reset_thread, total_reset_threads - ) - self._device_workspace.reset_tail_space( - "local", global_reset_thread, total_reset_threads - ) + self._device_workspace.reset_tail_space("shared", global_reset_thread, total_reset_threads) + self._device_workspace.reset_tail_space("local", global_reset_thread, total_reset_threads) iket.range_pop() iket.range_push("tail.nvlink_publish") @@ -1300,9 +1237,7 @@ def _kernel( iket.range_pop() iket.range_pop() - if (warp_idx >= Int32(self.gather_b_warp_ids[0])) & ( - warp_idx <= Int32(self.gather_b_warp_ids[-1]) - ): + if (warp_idx >= Int32(self.gather_b_warp_ids[0])) & (warp_idx <= Int32(self.gather_b_warp_ids[-1])): cute.arch.warpgroup_reg_dealloc(self.other_warp_register_count) iket.range_push("genphase.gather_b1") iket.range_push("gather_b1.wait_input_ready") diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/dynamic_mainloop.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/dynamic_mainloop.py index 56e119bbab85..4041a1fda3d8 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/dynamic_mainloop.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/dynamic_mainloop.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Raw SM107 dynamic-N block-scaled MMA emission for Rubin MegaMoE.""" from typing import Optional @@ -12,6 +9,7 @@ from .....quant_def import QuantKind + _bit_k_size_upper = 3 _bit_b_sf_id = 4 _bit_a_format = 7 @@ -56,9 +54,7 @@ def build_static_idesc_base( raise ValueError(f"SM107 MegaMoE requires instruction M 128 or 256, got {instruction_m}.") instruction_n_granularity = 16 if instruction_m == 256 else 8 if instruction_n <= 0 or instruction_n > 256 or instruction_n % instruction_n_granularity != 0: - raise ValueError( - f"Invalid SM107 instruction N {instruction_n} for instruction M {instruction_m}." - ) + raise ValueError(f"Invalid SM107 instruction N {instruction_n} for instruction M {instruction_m}.") expected_instruction_k = quant_kind.instruction_k("2x") if instruction_k != expected_instruction_k: @@ -146,15 +142,7 @@ def _tcgen05_mma_block_scaled( llvm.inline_asm( None, - [ - d_tmem_i32, - a_desc_i64, - b_desc_i64, - idesc_i32, - enable_input_d_i32, - sfa_tmem_i32, - sfb_tmem_i32, - ], + [d_tmem_i32, a_desc_i64, b_desc_i64, idesc_i32, enable_input_d_i32, sfa_tmem_i32, sfb_tmem_i32], "{\n\t" ".reg .pred p;\n\t" "setp.ne.b32 p, $4, 0;\n\t" diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/tma_gather.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/tma_gather.py new file mode 100644 index 000000000000..37d164b4a6e7 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/tma_gather.py @@ -0,0 +1,7 @@ +"""Rubin Genphase MegaMoE shim for the shared Gather4 primitive.""" + +# <<>> +from ..local_mega.tma_gather import sm107_tma_gather4_load + + +__all__ = ["sm107_tma_gather4_load"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/token_comm_gen_specialized.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/token_comm_gen_specialized.py index 99b5f26b170e..5adf1c789837 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/token_comm_gen_specialized.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/token_comm_gen_specialized.py @@ -47,9 +47,10 @@ """ import dataclasses -from typing import ClassVar, Optional +from typing import ClassVar, Optional, Tuple import cuda.bindings.driver as cuda + import cutlass import cutlass.cute as cute import cutlass.utils @@ -310,22 +311,16 @@ def _validate_configuration(self) -> None: if self.expert_count > 16384: raise NotImplementedError("Genphase token comm supports at most 16384 global experts.") if self.topk > self.expert_count: - raise ValueError( - f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}." - ) + raise ValueError(f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}.") if self.max_tokens_per_rank > 1024: raise ValueError( f"Genphase token comm requires max_tokens_per_rank <= 1024, got {self.max_tokens_per_rank}." ) if self.topk_index_dtype not in (cutlass.Int32, cutlass.Int64): - raise ValueError( - f"topk_index_dtype must be Int32 or Int64, got {self.topk_index_dtype}." - ) + raise ValueError(f"topk_index_dtype must be Int32 or Int64, got {self.topk_index_dtype}.") if self.sf_padding_block != 128: - raise ValueError( - f"sf_padding_block must be the 128-row SF atom, got {self.sf_padding_block}." - ) + raise ValueError(f"sf_padding_block must be the 128-row SF atom, got {self.sf_padding_block}.") if self.sf_padding_block % self.token_padding_block != 0: raise ValueError("sf_padding_block must be divisible by token_padding_block.") @@ -333,9 +328,7 @@ def _validate_configuration(self) -> None: # four of them along K, so the hidden extent has to fill whole atoms. elements_per_sf_atom = self.quant_kind.sf_vec_size * 4 if self.hidden_size % elements_per_sf_atom != 0: - raise ValueError( - f"{self.quant_kind} requires hidden_size divisible by {elements_per_sf_atom}." - ) + raise ValueError(f"{self.quant_kind} requires hidden_size divisible by {elements_per_sf_atom}.") # The caller's top-k plane is exactly `routes_per_rank` elements with no # slack, so the bulk copy that stages it has to land on that boundary @@ -363,9 +356,7 @@ def _validate_configuration(self) -> None: ) # One thread stages one row of a 128-row block. if self.refine_warps_per_group * 32 != 128: - raise ValueError( - "A refine group must be exactly four warps so one thread owns one row." - ) + raise ValueError("A refine group must be exactly four warps so one thread owns one row.") # The staging load reads a run-aligned row straight out of the linear # plane, so the plane's row has to be at least that wide. Both extents # round the same value up to sixteen, from a payload alignment and from an @@ -540,9 +531,7 @@ def worst_case_padded_rows(self, block: int) -> int: routes_per_source_token = min(self.topk, self.experts_per_rank) route_capacity = source_token_capacity * routes_per_source_token active_expert_capacity = min(self.experts_per_rank, route_capacity) - route_budget_blocks = ( - active_expert_capacity + (route_capacity - active_expert_capacity) // block - ) + route_budget_blocks = active_expert_capacity + (route_capacity - active_expert_capacity) // block expert_bound_blocks = active_expert_capacity * int(ceil_div(source_token_capacity, block)) return min(route_budget_blocks, expert_bound_blocks) * block @@ -634,23 +623,14 @@ def _register_symmetric_regions(self, workspace: DeviceWorkspace) -> None: byte_alignment=16, ) workspace.register( - self.topk_ready_region, - cutlass.Int32, - (self.world_size,), - buffer_space="shared", - reset="tail_reset", + self.topk_ready_region, cutlass.Int32, (self.world_size,), buffer_space="shared", reset="tail_reset" ) # The shape is the logical column extent and the row padding is a stride. # Bulk copies move whole padded rows, so the trailing row's padding has to # be inside the region even though no index reaches it; the workspace sizes # a region by what its strides occupy for exactly this reason. for region_name, dtype, columns, row_elements in ( - ( - self.plain_activation_region, - self.activation_dtype, - self.hidden_size, - self.plain_activation_row_elements, - ), + (self.plain_activation_region, self.activation_dtype, self.hidden_size, self.plain_activation_row_elements), ( self.plain_activation_sf_region, self.activation_sf_dtype, @@ -667,9 +647,7 @@ def _register_symmetric_regions(self, workspace: DeviceWorkspace) -> None: byte_alignment=128, ) for region_name in (self.token_data_ready_region, self.token_sf_plain_ready_region): - workspace.register( - region_name, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset" - ) + workspace.register(region_name, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") def _register_local_regions(self, workspace: DeviceWorkspace) -> None: workspace.register( @@ -681,11 +659,7 @@ def _register_local_regions(self, workspace: DeviceWorkspace) -> None: byte_alignment=16, ) workspace.register( - self.expert_sizes_region, - cutlass.Int32, - (self.experts_per_rank,), - buffer_space="local", - reset="tail_reset", + self.expert_sizes_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local", reset="tail_reset" ) workspace.register( self.source_expert_base_region, @@ -712,17 +686,9 @@ def _register_local_regions(self, workspace: DeviceWorkspace) -> None: # Sized to the padded bucket count rather than the live expert count: the # scan that fills them runs over padded buckets, and letting the tail # entries exist removes every boundary special case downstream. - for region_name in ( - self.data_expert_base_region, - self.sf_expert_base_region, - self.refine_block_base_region, - ): + for region_name in (self.data_expert_base_region, self.sf_expert_base_region, self.refine_block_base_region): workspace.register( - region_name, - cutlass.Int32, - (self.local_expert_bucket_count,), - buffer_space="local", - byte_alignment=16, + region_name, cutlass.Int32, (self.local_expert_bucket_count,), buffer_space="local", byte_alignment=16 ) workspace.register( self.token_src_metadata_region, @@ -771,16 +737,10 @@ def _register_local_regions(self, workspace: DeviceWorkspace) -> None: self.refine_block_total_region, self.refine_work_counter_region, ): - workspace.register( - region_name, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" - ) + workspace.register(region_name, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset") # One slot per local expert workspace.register( - self.fc1_sf_ready_region, - cutlass.Int32, - (self.experts_per_rank,), - buffer_space="local", - reset="tail_reset", + self.fc1_sf_ready_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local", reset="tail_reset" ) # ------------------------------------------------------------------ @@ -798,9 +758,7 @@ def _build_smem_workspace(self) -> SmemWorkspace: warp_count = self.threads_per_cta // 32 workspace = SmemWorkspace() workspace.register_mbarrier(self.exchange_mbarrier_region, 1) - workspace.register_mbarrier( - self.pusher_mbarrier_region, warp_count * self.pusher_pipeline_depth - ) + workspace.register_mbarrier(self.pusher_mbarrier_region, warp_count * self.pusher_pipeline_depth) workspace.register_mbarrier(self.refine_mbarrier_region, self.refine_participant_groups) overlay = workspace.create_overlay("genphase.token_comm.smem.role") @@ -810,37 +768,22 @@ def _build_smem_workspace(self) -> SmemWorkspace: # what the transfer reads. exchange = overlay.add_lifetime("exchange") exchange.register_tensor( - self.exchange_topk_stage_region, - self.topk_index_dtype, - (self.topk_slab_elements,), - byte_alignment=16, + self.exchange_topk_stage_region, self.topk_index_dtype, (self.topk_slab_elements,), byte_alignment=16 ) if self.apply_topk_at_fc1: exchange.register_tensor( - self.exchange_score_stage_region, - cutlass.Float32, - (self.topk_slab_elements,), - byte_alignment=16, + self.exchange_score_stage_region, cutlass.Float32, (self.topk_slab_elements,), byte_alignment=16 ) sort = overlay.add_lifetime("sort") sort.register_tensor( - self.sort_histogram_region, - cutlass.Int32, - (self.histogram_bucket_count,), - byte_alignment=16, + self.sort_histogram_region, cutlass.Int32, (self.histogram_bucket_count,), byte_alignment=16 ) sort.register_tensor( - self.sort_prefix_region, - cutlass.Int32, - (self.local_expert_bucket_count,), - byte_alignment=16, + self.sort_prefix_region, cutlass.Int32, (self.local_expert_bucket_count,), byte_alignment=16 ) sort.register_tensor( - self.sort_warp_totals_region, - cutlass.Int32, - (self.threads_per_cta // 32,), - byte_alignment=16, + self.sort_warp_totals_region, cutlass.Int32, (self.threads_per_cta // 32,), byte_alignment=16 ) helper = overlay.add_lifetime("helper") @@ -851,44 +794,22 @@ def _build_smem_workspace(self) -> SmemWorkspace: stride=(self.local_expert_bucket_count, 1), byte_alignment=16, ) - for region_name in ( - self.helper_totals_region, - self.helper_scan_input_region, - self.helper_prefix_region, - ): - helper.register_tensor( - region_name, cutlass.Int32, (self.local_expert_bucket_count,), byte_alignment=16 - ) + for region_name in (self.helper_totals_region, self.helper_scan_input_region, self.helper_prefix_region): + helper.register_tensor(region_name, cutlass.Int32, (self.local_expert_bucket_count,), byte_alignment=16) helper.register_tensor( - self.helper_warp_totals_region, - cutlass.Int32, - (self.threads_per_cta // 32,), - byte_alignment=16, + self.helper_warp_totals_region, cutlass.Int32, (self.threads_per_cta // 32,), byte_alignment=16 ) for region_name in (self.helper_slot_to_expert_region, self.helper_totals_by_slot_region): - helper.register_tensor( - region_name, cutlass.Int32, (self.local_expert_bucket_count,), byte_alignment=16 - ) + helper.register_tensor(region_name, cutlass.Int32, (self.local_expert_bucket_count,), byte_alignment=16) if not self.ranks_experts_in_warp: helper.register_tensor( - self.helper_load_bin_region, - cutlass.Int32, - (self.load_bin_count,), - byte_alignment=16, + self.helper_load_bin_region, cutlass.Int32, (self.load_bin_count,), byte_alignment=16 ) pusher = overlay.add_lifetime("pusher") for region_name, dtype, row_elements in ( - ( - self.pusher_data_stage_region, - self.activation_dtype, - self.plain_activation_row_elements, - ), - ( - self.pusher_sf_stage_region, - self.activation_sf_dtype, - self.plain_activation_sf_row_elements, - ), + (self.pusher_data_stage_region, self.activation_dtype, self.plain_activation_row_elements), + (self.pusher_sf_stage_region, self.activation_sf_dtype, self.plain_activation_sf_row_elements), ): pusher.register_tensor( region_name, @@ -926,11 +847,7 @@ def _build_smem_workspace(self) -> SmemWorkspace: refine.register_tensor( self.refine_output_region, self.activation_sf_dtype, - ( - self.refine_participant_groups, - self.refine_output_stages, - self.refine_output_stage_elements, - ), + (self.refine_participant_groups, self.refine_output_stages, self.refine_output_stage_elements), stride=( self.refine_output_stages * self.refine_output_stage_elements, self.refine_output_stage_elements, @@ -965,9 +882,7 @@ def topk_score_inbox_tensor(self, device_workspace: DeviceWorkspace) -> Optional return device_workspace.tensor(self.topk_score_inbox_region) @cute.jit - def _payload_plane_by_dense_row( - self, device_workspace: DeviceWorkspace, region_name: str - ) -> cute.Tensor: + def _payload_plane_by_dense_row(self, device_workspace: DeviceWorkspace, region_name: str) -> cute.Tensor: """``(world * tokens, columns)`` view of a payload plane. The registered shape keeps the source rank as its own mode, which is how @@ -982,9 +897,7 @@ def _payload_plane_by_dense_row( row_elements = device_workspace.stride(region_name)[1] return cute.make_tensor( device_workspace.ptr(region_name), - cute.make_layout( - (self.world_size * self.max_tokens_per_rank, columns), stride=(row_elements, 1) - ), + cute.make_layout((self.world_size * self.max_tokens_per_rank, columns), stride=(row_elements, 1)), ) @cute.jit @@ -1010,12 +923,9 @@ def canonical_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> c ``tile_atom_to_shape_SF`` takes the *data* shape, not the SF shape: the broadcast mode it inserts already absorbs ``sf_vec_size``. """ - layout = tile_atom_to_shape_SF( - (self.worst_case_sf_rows, self.hidden_size, 1), self.quant_kind.sf_vec_size - ) + layout = tile_atom_to_shape_SF((self.worst_case_sf_rows, self.hidden_size, 1), self.quant_kind.sf_vec_size) return cute.make_tensor( - device_workspace.ptr(self.canonical_activation_sf_region), - cute.select(layout, mode=[0, 1]), + device_workspace.ptr(self.canonical_activation_sf_region), cute.select(layout, mode=[0, 1]) ) @cute.jit @@ -1077,13 +987,9 @@ def launch( f"topk_indices must be the declared {self.topk_index_dtype}, got {args.topk_indices.element_type}." ) if cutlass.const_expr(args.activation.element_type is not self.activation_dtype): - raise TypeError( - f"activation must be {self.activation_dtype}, got {args.activation.element_type}." - ) + raise TypeError(f"activation must be {self.activation_dtype}, got {args.activation.element_type}.") if cutlass.const_expr(args.activation_sf.element_type is not self.activation_sf_dtype): - raise TypeError( - f"activation_sf must be {self.activation_sf_dtype}, got {args.activation_sf.element_type}." - ) + raise TypeError(f"activation_sf must be {self.activation_sf_dtype}, got {args.activation_sf.element_type}.") # Row addressing uses each plane's own layout, so the padded stride is a # requirement rather than something this component can paper over: a whole # padded row is bulk-copied, and a tighter caller stride would be read past @@ -1095,9 +1001,7 @@ def launch( ): if cutlass.const_expr(isinstance(plane.stride[0], int)): if cutlass.const_expr(plane.stride[0] != row_elements): - raise ValueError( - f"{plane_name} needs the padded row stride {row_elements}, got {plane.stride[0]}." - ) + raise ValueError(f"{plane_name} needs the padded row stride {row_elements}, got {plane.stride[0]}.") self._kernel( args.topk_indices, @@ -1234,19 +1138,10 @@ def _run_exchange( if thread_idx == Int32(0): cute.arch.mbarrier_arrive_and_expect_tx(load_mbarrier, Int32(staged_bytes)) - tma_load_1d( - index_stage.iterator, topk_indices.iterator, load_mbarrier, Int32(index_plane_bytes) - ) + tma_load_1d(index_stage.iterator, topk_indices.iterator, load_mbarrier, Int32(index_plane_bytes)) if cutlass.const_expr(self.apply_topk_at_fc1): - score_stage = self._smem_workspace.tensor( - self.exchange_score_stage_region, smem_base - ) - tma_load_1d( - score_stage.iterator, - topk_scores.iterator, - load_mbarrier, - Int32(score_plane_bytes), - ) + score_stage = self._smem_workspace.tensor(self.exchange_score_stage_region, smem_base) + tma_load_1d(score_stage.iterator, topk_scores.iterator, load_mbarrier, Int32(score_plane_bytes)) cute.arch.mbarrier_wait(load_mbarrier, 0) self._canonicalize_wire_indices(index_stage, thread_idx) @@ -1265,10 +1160,7 @@ def _run_exchange( if thread_idx == Int32(0): cp_async_bulk_s2g( cute.make_ptr( - cutlass.Int32, - own_index_slab.iterator.toint() + peer_offset, - AddressSpace.gmem, - assumed_align=16, + cutlass.Int32, own_index_slab.iterator.toint() + peer_offset, AddressSpace.gmem, assumed_align=16 ), wire_indices.iterator, Int32(self.topk_slab_elements * 4), @@ -1283,9 +1175,7 @@ def _run_exchange( AddressSpace.gmem, assumed_align=16, ), - self._smem_workspace.tensor( - self.exchange_score_stage_region, smem_base - ).iterator, + self._smem_workspace.tensor(self.exchange_score_stage_region, smem_base).iterator, Int32(self.topk_slab_elements * 4), ) cute.arch.cp_async_bulk_commit_group() @@ -1311,8 +1201,7 @@ def _wire_index_view(self, index_stage: cute.Tensor) -> cute.Tensor: if cutlass.const_expr(self.topk_index_dtype is cutlass.Int32): return index_stage return cute.make_tensor( - cute.recast_ptr(index_stage.iterator, dtype=cutlass.Int32), - cute.make_layout((self.topk_slab_elements,)), + cute.recast_ptr(index_stage.iterator, dtype=cutlass.Int32), cute.make_layout((self.topk_slab_elements,)) ) @cute.jit @@ -1345,9 +1234,7 @@ def _is_unowned_expert(self, raw_index) -> cutlass.Boolean: return cutlass.Uint64(raw_index) >= cutlass.Uint64(self.expert_count) @cute.jit - def _classify_tokens( - self, topk_indices: cute.Tensor, global_warp_idx: Int32, lane_idx: Int32 - ) -> cute.Tensor: + def _classify_tokens(self, topk_indices: cute.Tensor, global_warp_idx: Int32, lane_idx: Int32) -> cute.Tensor: """One bit per (token, peer): does this warp owe that peer this token? Both payload passes need the same answer from the same input, so deriving @@ -1388,13 +1275,9 @@ def _classify_tokens( if lane_owns_slot: raw_index = topk_indices[token, slot] if not self._is_unowned_expert(raw_index): - lane_rank = cutlass.Uint32(raw_index) // cutlass.Uint32( - self.experts_per_rank - ) + lane_rank = cutlass.Uint32(raw_index) // cutlass.Uint32(self.experts_per_rank) if cutlass.const_expr(mask_words == 1): - masks[token_round, 0] = masks[token_round, 0] | ( - cutlass.Uint32(1) << lane_rank - ) + masks[token_round, 0] = masks[token_round, 0] | (cutlass.Uint32(1) << lane_rank) else: lane_word = lane_rank // cutlass.Uint32(32) lane_bit = cutlass.Uint32(1) << (lane_rank % cutlass.Uint32(32)) @@ -1404,9 +1287,7 @@ def _classify_tokens( for token_round in cutlass.range_constexpr(tokens_per_warp): for word in cutlass.range_constexpr(mask_words): - masks[token_round, word] = cutlass.Uint32( - cute.arch.warp_redux_sync(masks[token_round, word], "or") - ) + masks[token_round, word] = cutlass.Uint32(cute.arch.warp_redux_sync(masks[token_round, word], "or")) return masks @cute.jit @@ -1451,9 +1332,7 @@ def _run_pusher( iket.range_push("pusher.classify") destination_masks = self._classify_tokens(topk_indices, global_warp_idx, lane_idx) iket.range_pop() - stage_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.pusher_pipeline_depth - ) + stage_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.pusher_pipeline_depth) iket.range_push("pusher.sf") stage_state = self._push_one_plane( @@ -1538,10 +1417,7 @@ def _push_one_plane( with cute.arch.elect_one(): cute.arch.mbarrier_arrive_and_expect_tx(slot_mbarrier, Int32(row_bytes)) tma_load_1d( - slot_stage.iterator, - source_plane[token, None].iterator, - slot_mbarrier, - Int32(row_bytes), + slot_stage.iterator, source_plane[token, None].iterator, slot_mbarrier, Int32(row_bytes) ) cute.arch.mbarrier_wait(slot_mbarrier, stage_state.phase) iket.range_pop() @@ -1550,9 +1426,7 @@ def _push_one_plane( for destination_rank in cutlass.range_constexpr(self.world_size): mask_word = destination_rank // 32 peer_bit = cutlass.Uint32(1 << (destination_rank % 32)) - peer_offset = peer_rank_ptr_mapper.map( - Int64(0), Int32(destination_rank), Int64(0) - ) + peer_offset = peer_rank_ptr_mapper.map(Int64(0), Int32(destination_rank), Int64(0)) destination_row = destination_plane[local_rank, token, None] if (destination_masks[token_round, mask_word] & peer_bit) != cutlass.Uint32(0): with cute.arch.elect_one(): @@ -1575,10 +1449,7 @@ def _push_one_plane( @cute.jit def _publish_payload_flags( - self, - device_workspace: DeviceWorkspace, - peer_rank_ptr_mapper: SymmetricBufferDevice, - thread_idx: Int32, + self, device_workspace: DeviceWorkspace, peer_rank_ptr_mapper: SymmetricBufferDevice, thread_idx: Int32 ) -> None: """Drain both payload passes, then tell every peer this pusher is finished. @@ -1599,9 +1470,7 @@ def _publish_payload_flags( if thread_idx < Int32(self.world_size): for region_name in (self.token_sf_plain_ready_region, self.token_data_ready_region): counter_address = device_workspace.ptr(region_name).toint() - red_add_relaxed_sys_s32( - peer_rank_ptr_mapper.map(counter_address, thread_idx, Int64(0)), Int32(1) - ) + red_add_relaxed_sys_s32(peer_rank_ptr_mapper.map(counter_address, thread_idx, Int64(0)), Int32(1)) @cute.jit def _run_sort( @@ -1624,9 +1493,7 @@ def _run_sort( histogram = self._smem_workspace.tensor(self.sort_histogram_region, smem_base) base_row = self._smem_workspace.tensor(self.sort_prefix_region, smem_base) - for bucket_round in cutlass.range_constexpr( - ceil_div(self.histogram_bucket_count, self.threads_per_cta) - ): + for bucket_round in cutlass.range_constexpr(ceil_div(self.histogram_bucket_count, self.threads_per_cta)): bucket = Int32(bucket_round * self.threads_per_cta) + thread_idx if bucket < Int32(self.histogram_bucket_count): histogram[bucket] = Int32(0) @@ -1661,25 +1528,18 @@ def _run_sort( if route < Int32(self.routes_per_rank): # One unsigned compare covers "not this rank's" and "invalid": # the wire marker and any wild value both land outside the window. - local_expert = ( - index_slab[route // Int32(self.topk), route % Int32(self.topk)] - - local_expert_begin - ) + local_expert = index_slab[route // Int32(self.topk), route % Int32(self.topk)] - local_expert_begin if cutlass.Uint32(local_expert) < cutlass.Uint32(self.experts_per_rank): buckets[route_round] = local_expert for route_round in cutlass.range_constexpr(route_rounds): offsets[route_round] = Int32( - cute.arch.atomic_add( - histogram.iterator + buckets[route_round], Int32(1), sem="relaxed", scope="cta" - ) + cute.arch.atomic_add(histogram.iterator + buckets[route_round], Int32(1), sem="relaxed", scope="cta") ) cute.arch.sync_threads() expert_counts = device_workspace.tensor(self.expert_counts_by_rank_region) - for bucket_round in cutlass.range_constexpr( - ceil_div(self.local_expert_bucket_count, self.threads_per_cta) - ): + for bucket_round in cutlass.range_constexpr(ceil_div(self.local_expert_bucket_count, self.threads_per_cta)): bucket = Int32(bucket_round * self.threads_per_cta) + thread_idx if bucket < Int32(self.local_expert_bucket_count): expert_counts[source_rank, bucket] = histogram[bucket] @@ -1688,14 +1548,8 @@ def _run_sort( iket.range_push("sort.wait_prefix") if thread_idx == Int32(0): - red_async_add_release_gpu_s32( - device_workspace.ptr(self.rank_count_ready_region), Int32(1) - ) - spin_wait( - device_workspace.ptr(self.prefix_ready_region), - lambda value: value == Int32(1), - scope="gpu", - ) + red_async_add_release_gpu_s32(device_workspace.ptr(self.rank_count_ready_region), Int32(1)) + spin_wait(device_workspace.ptr(self.prefix_ready_region), lambda value: value == Int32(1), scope="gpu") cute.arch.sync_threads() iket.range_pop() @@ -1703,9 +1557,7 @@ def _run_sort( # This source's segment start inside each expert, staged once: every thread # needs it for up to `route_rounds` unrelated buckets. source_expert_base = device_workspace.tensor(self.source_expert_base_region) - for bucket_round in cutlass.range_constexpr( - ceil_div(self.local_expert_bucket_count, self.threads_per_cta) - ): + for bucket_round in cutlass.range_constexpr(ceil_div(self.local_expert_bucket_count, self.threads_per_cta)): bucket = Int32(bucket_round * self.threads_per_cta) + thread_idx if bucket < Int32(self.local_expert_bucket_count): base_row[bucket] = source_expert_base[source_rank, bucket] @@ -1715,9 +1567,7 @@ def _run_sort( token_metadata = device_workspace.tensor(self.token_src_metadata_region) routed_scores = self.fc1_topk_scores_tensor(device_workspace) if cutlass.const_expr(self.apply_topk_at_fc1): - score_slab = device_workspace.tensor(self.topk_score_inbox_region)[ - source_rank, None, None - ] + score_slab = device_workspace.tensor(self.topk_score_inbox_region)[source_rank, None, None] dense_row_begin = source_rank * Int32(self.max_tokens_per_rank) for route_round in cutlass.range_constexpr(route_rounds): @@ -1735,18 +1585,12 @@ def _run_sort( cute.arch.sync_threads() if thread_idx == Int32(0): - red_async_add_release_gpu_s32( - device_workspace.ptr(self.metadata_ready_region), Int32(1) - ) + red_async_add_release_gpu_s32(device_workspace.ptr(self.metadata_ready_region), Int32(1)) iket.range_pop() @cute.jit def _run_helper( - self, - device_workspace: DeviceWorkspace, - local_rank: Int32, - smem_base: cute.Pointer, - thread_idx: Int32, + self, device_workspace: DeviceWorkspace, local_rank: Int32, smem_base: cute.Pointer, thread_idx: Int32 ) -> None: """Fold the W count rows into per-expert sizes and both row-space prefixes. @@ -1814,12 +1658,8 @@ def _run_helper( # themselves stay keyed by expert, because every reader of them -- the # source segment bases, the refine expert table, and the main kernel's # gather -- addresses an expert rather than a slot. - loads_in_scan_order = self._smem_workspace.tensor( - self.helper_totals_by_slot_region, smem_base - ) - scan_order_experts = self._smem_workspace.tensor( - self.helper_slot_to_expert_region, smem_base - ) + loads_in_scan_order = self._smem_workspace.tensor(self.helper_totals_by_slot_region, smem_base) + scan_order_experts = self._smem_workspace.tensor(self.helper_slot_to_expert_region, smem_base) for padding_block, region_name in ( (self.token_padding_block, self.data_expert_base_region), @@ -1881,9 +1721,7 @@ def _rank_experts_by_load( if cutlass.const_expr(self.ranks_experts_in_warp): self._rank_by_counting_in_warp(device_workspace, smem_base, lane_idx, warp_idx) else: - self._rank_by_load_histogram( - device_workspace, smem_base, thread_idx, lane_idx, warp_idx - ) + self._rank_by_load_histogram(device_workspace, smem_base, thread_idx, lane_idx, warp_idx) @cute.jit def _place_expert_at_slot( @@ -1907,20 +1745,14 @@ def _place_expert_at_slot( kernel never translates it. """ self._smem_workspace.tensor(self.helper_slot_to_expert_region, smem_base)[slot] = expert - self._smem_workspace.tensor(self.helper_totals_by_slot_region, smem_base)[slot] = ( - routed_count - ) + self._smem_workspace.tensor(self.helper_totals_by_slot_region, smem_base)[slot] = routed_count device_workspace.tensor(self.slot_to_expert_region)[slot] = expert if slot < Int32(self.experts_per_rank): device_workspace.tensor(self.expert_sizes_by_slot_region)[slot] = routed_count @cute.jit def _rank_by_counting_in_warp( - self, - device_workspace: DeviceWorkspace, - smem_base: cute.Pointer, - lane_idx: Int32, - warp_idx: Int32, + self, device_workspace: DeviceWorkspace, smem_base: cute.Pointer, lane_idx: Int32, warp_idx: Int32 ) -> None: """Rank the experts inside one warp by counting how many outrank each. @@ -1993,12 +1825,7 @@ def _rank_by_load_histogram( binade = Int32(32) - Int32(cute.arch.clz(totals[bucket])) keys[bucket_round] = Int32(self.load_bin_count - 1) - binade offsets_in_bin[bucket_round] = Int32( - cute.arch.atomic_add( - load_bins.iterator + keys[bucket_round], - Int32(1), - sem="relaxed", - scope="cta", - ) + cute.arch.atomic_add(load_bins.iterator + keys[bucket_round], Int32(1), sem="relaxed", scope="cta") ) cute.arch.sync_threads() @@ -2015,9 +1842,7 @@ def _rank_by_load_histogram( bucket = Int32(bucket_round * self.threads_per_cta) + thread_idx if bucket < Int32(self.local_expert_bucket_count): slot = load_bins[keys[bucket_round]] + offsets_in_bin[bucket_round] - self._place_expert_at_slot( - device_workspace, smem_base, slot, bucket, totals[bucket] - ) + self._place_expert_at_slot(device_workspace, smem_base, slot, bucket, totals[bucket]) cute.arch.sync_threads() @cute.jit @@ -2051,14 +1876,10 @@ def _publish_padded_prefix( for bucket_round in cutlass.range_constexpr(bucket_rounds): scan_entry = Int32(bucket_round * self.threads_per_cta) + thread_idx if scan_entry < Int32(self.local_expert_bucket_count): - scan_input[scan_entry] = padded_expert_rows( - loads_in_scan_order[scan_entry], padding_block - ) + scan_input[scan_entry] = padded_expert_rows(loads_in_scan_order[scan_entry], padding_block) cute.arch.sync_threads() - smem_exclusive_prefix( - scan_input, prefix, warp_totals, self.threads_per_cta, thread_idx, lane_idx, warp_idx - ) + smem_exclusive_prefix(scan_input, prefix, warp_totals, self.threads_per_cta, thread_idx, lane_idx, warp_idx) cute.arch.sync_threads() destination = device_workspace.tensor(region_name) @@ -2100,9 +1921,7 @@ def _publish_refine_block_prefix( for bucket_round in cutlass.range_constexpr(bucket_rounds): scan_entry = Int32(bucket_round * self.threads_per_cta) + thread_idx if scan_entry < Int32(self.local_expert_bucket_count): - scan_input[scan_entry] = ceil_div( - loads_in_scan_order[scan_entry], Int32(self.sf_padding_block) - ) + scan_input[scan_entry] = ceil_div(loads_in_scan_order[scan_entry], Int32(self.sf_padding_block)) cute.arch.sync_threads() block_total = smem_exclusive_prefix( @@ -2121,11 +1940,7 @@ def _publish_refine_block_prefix( @cute.jit def _publish_source_segment_bases( - self, - device_workspace: DeviceWorkspace, - count_matrix: cute.Tensor, - local_rank: Int32, - thread_idx: Int32, + self, device_workspace: DeviceWorkspace, count_matrix: cute.Tensor, local_rank: Int32, thread_idx: Int32 ) -> None: """Where each (source rank, expert) segment starts in the data row space. @@ -2136,9 +1951,7 @@ def _publish_source_segment_bases( data_expert_base = device_workspace.tensor(self.data_expert_base_region) source_expert_base = device_workspace.tensor(self.source_expert_base_region) - for bucket_round in cutlass.range_constexpr( - ceil_div(self.local_expert_bucket_count, self.threads_per_cta) - ): + for bucket_round in cutlass.range_constexpr(ceil_div(self.local_expert_bucket_count, self.threads_per_cta)): bucket = Int32(bucket_round * self.threads_per_cta) + thread_idx if bucket < Int32(self.local_expert_bucket_count): segment_begin = data_expert_base[bucket] @@ -2149,9 +1962,7 @@ def _publish_source_segment_bases( cute.arch.sync_threads() @cute.jit - def _run_refine( - self, device_workspace: DeviceWorkspace, smem_base: cute.Pointer, thread_idx: Int32 - ) -> None: + def _run_refine(self, device_workspace: DeviceWorkspace, smem_base: cute.Pointer, thread_idx: Int32) -> None: """Rewrite the received scale factors into the layout the MMA reads.""" group_threads = self.refine_warps_per_group * 32 group_idx = thread_idx // Int32(group_threads) @@ -2189,21 +2000,12 @@ def _run_refine( work_counter = device_workspace.ptr(self.refine_work_counter_region) block_total = device_workspace.tensor(self.refine_block_total_region)[0] - claim_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=128 - ) + claim_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=128) claimed = cute.make_rmem_tensor((self.refine_claim_fields,), cutlass.Int32) fc1_sf_ready = device_workspace.ptr(self.fc1_sf_ready_region) owning_expert = self._claim_next_block( - claim, - expert_table, - work_counter, - block_total, - group_idx, - thread_in_group, - group_barrier, - group_threads, + claim, expert_table, work_counter, block_total, group_idx, thread_in_group, group_barrier, group_threads ) cute.copy(claim_atom, claim[group_idx, None], claimed) @@ -2227,14 +2029,7 @@ def _run_refine( red_async_add_release_gpu_s32(fc1_sf_ready + owning_expert, Int32(1)) iket.range_pop() owning_expert = self._claim_next_block( - claim, - expert_table, - work_counter, - block_total, - group_idx, - thread_in_group, - group_barrier, - group_threads, + claim, expert_table, work_counter, block_total, group_idx, thread_in_group, group_barrier, group_threads ) cute.copy(claim_atom, claim[group_idx, None], claimed) @@ -2269,9 +2064,7 @@ def _claim_next_block( if thread_in_group < Int32(32): claimed_block = Int32(0) if thread_in_group == Int32(0): - claimed_block = Int32( - cute.arch.atomic_add(work_counter, Int32(1), sem="relaxed", scope="gpu") - ) + claimed_block = Int32(cute.arch.atomic_add(work_counter, Int32(1), sem="relaxed", scope="gpu")) # Broadcast rather than staged through shared memory: the ballot below # needs the whole warp to agree on the block, and a shuffle gets it # there without a second barrier. @@ -2284,25 +2077,17 @@ def _claim_next_block( resolved[self.refine_claim_live_rows] = Int32(0) if claimed_block < block_total: block_rows = Int32(self.sf_padding_block) - owning_expert = self._resolve_block_expert( - expert_table, claimed_block, thread_in_group - ) - rows_before = ( - claimed_block - expert_table[self.refine_table_block_base, owning_expert] - ) * block_rows + owning_expert = self._resolve_block_expert(expert_table, claimed_block, thread_in_group) + rows_before = (claimed_block - expert_table[self.refine_table_block_base, owning_expert]) * block_rows expert_size = expert_table[self.refine_table_expert_size, owning_expert] data_base = expert_table[self.refine_table_data_base, owning_expert] sf_base = expert_table[self.refine_table_sf_base, owning_expert] resolved[self.refine_claim_data_row] = data_base + rows_before resolved[self.refine_claim_sf_row] = sf_base + rows_before - resolved[self.refine_claim_live_rows] = cutlass.min( - expert_size - rows_before, block_rows - ) + resolved[self.refine_claim_live_rows] = cutlass.min(expert_size - rows_before, block_rows) if thread_in_group == Int32(0): - claim_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=128 - ) + claim_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=128) cute.copy(claim_atom, resolved, claim[group_idx, None]) cute.arch.barrier(barrier_id=group_barrier, number_of_threads=group_threads) return owning_expert @@ -2322,9 +2107,7 @@ def _stage_refine_expert_table( sf_base = device_workspace.tensor(self.sf_expert_base_region) expert_sizes = device_workspace.tensor(self.expert_sizes_region) - for bucket_round in cutlass.range_constexpr( - ceil_div(self.local_expert_bucket_count, self.threads_per_cta) - ): + for bucket_round in cutlass.range_constexpr(ceil_div(self.local_expert_bucket_count, self.threads_per_cta)): bucket = Int32(bucket_round * self.threads_per_cta) + thread_idx if bucket < Int32(self.local_expert_bucket_count): expert_table[self.refine_table_block_base, bucket] = block_base[bucket] @@ -2377,21 +2160,13 @@ def _refine_one_block( iket.range_push("refine.out_of_smem") self._scatter_block_atoms( - device_workspace, - smem_base, - sf_row_begin, - group_idx, - thread_in_group, - group_barrier, - group_threads, + device_workspace, smem_base, sf_row_begin, group_idx, thread_in_group, group_barrier, group_threads ) iket.range_pop() return stage_state @cute.jit - def _resolve_block_expert( - self, expert_table: cute.Tensor, block_index: Int32, lane_idx: Int32 - ) -> Int32: + def _resolve_block_expert(self, expert_table: cute.Tensor, block_index: Int32, lane_idx: Int32) -> Int32: """Which local expert owns a linear refine block. One lane per bucket, so the answer comes out of a ballot rather than a scan @@ -2411,12 +2186,9 @@ def _resolve_block_expert( if bucket < Int32(self.local_expert_bucket_count): bucket_base = expert_table[self.refine_table_block_base, bucket] bucket_blocks = ceil_div( - expert_table[self.refine_table_expert_size, bucket], - Int32(self.sf_padding_block), - ) - bucket_owns_block = (bucket_base <= block_index) & ( - block_index < bucket_base + bucket_blocks + expert_table[self.refine_table_expert_size, bucket], Int32(self.sf_padding_block) ) + bucket_owns_block = (bucket_base <= block_index) & (block_index < bucket_base + bucket_blocks) owning_lane = _first_matching_lane(bucket_owns_block) if owning_lane >= Int32(0): resolved = Int32(bucket_round * 32) + owning_lane @@ -2507,51 +2279,32 @@ def _scatter_block_atoms( # The runs tile the row exactly: the staging tile's column extent is the # run-aligned one, so a partial last run reads columns past `hidden_sf` # that belong to this row rather than to the next. - staged_runs = cute.zipped_divide( - staging[group_idx, thread_in_group, None], (columns_per_run,) - ) + staged_runs = cute.zipped_divide(staging[group_idx, thread_in_group, None], (columns_per_run,)) load_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), - self.activation_sf_dtype, - num_bits_per_copy=8 * columns_per_run, - ) - store_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), self.activation_sf_dtype, num_bits_per_copy=32 - ) - staged_values = cute.make_rmem_tensor( - cute.make_layout((columns_per_run,)), self.activation_sf_dtype + cute.nvgpu.CopyUniversalOp(), self.activation_sf_dtype, num_bits_per_copy=8 * columns_per_run ) + store_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.activation_sf_dtype, num_bits_per_copy=32) + staged_values = cute.make_rmem_tensor(cute.make_layout((columns_per_run,)), self.activation_sf_dtype) value_atoms = cute.zipped_divide(staged_values, (4,)) for run in cutlass.range_constexpr(self.refine_run_count): stage_slot = run % self.refine_output_stages - live_atoms = min( - self.refine_atoms_per_stage, - self.refine_total_atoms - run * self.refine_atoms_per_stage, - ) + live_atoms = min(self.refine_atoms_per_stage, self.refine_total_atoms - run * self.refine_atoms_per_stage) # One run of atoms, addressed the way the MMA reads it: 32 rows by 4 # row-groups against 4 scale factors by 4 atoms. run_output = cute.make_tensor( output[group_idx, stage_slot, None].iterator, - cute.make_layout( - ((32, 4), (4, self.refine_atoms_per_stage)), stride=((16, 4), (1, 512)) - ), + cute.make_layout(((32, 4), (4, self.refine_atoms_per_stage)), stride=((16, 4), (1, 512))), ) cute.copy(load_atom, staged_runs[None, run], staged_values) for atom in cutlass.range_constexpr(self.refine_atoms_per_stage): - cute.copy( - store_atom, - value_atoms[None, atom], - run_output[(atom_row, atom_row_group), (None, atom)], - ) + cute.copy(store_atom, value_atoms[None, atom], run_output[(atom_row, atom_row_group), (None, atom)]) cute.arch.barrier(barrier_id=group_barrier, number_of_threads=group_threads) # `sf_row_begin` is atom-aligned and the run starts on an atom # boundary, so the tensor's own layout hands back the run base. - destination = cute.domain_offset( - (sf_row_begin, Int32(run * columns_per_run * sf_vec_size)), canonical - ) + destination = cute.domain_offset((sf_row_begin, Int32(run * columns_per_run * sf_vec_size)), canonical) # Elected inside the group's first warp, not in every warp: the stage # slot is written cooperatively by all `group_threads` threads, so it # leaves as one store. A bare `elect_one` would issue the same address @@ -2562,9 +2315,7 @@ def _scatter_block_atoms( if thread_in_group // Int32(32) == Int32(0): with cute.arch.elect_one(): cp_async_bulk_s2g( - destination.iterator, - output[group_idx, stage_slot, None].iterator, - Int32(live_atoms * 512), + destination.iterator, output[group_idx, stage_slot, None].iterator, Int32(live_atoms * 512) ) cute.arch.cp_async_bulk_commit_group() cute.arch.cp_async_bulk_wait_group(self.refine_output_stages - 1, read=True) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/base.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/base.py index d6e22636108e..f4a1f40efda1 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/base.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/base.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Scheduler façade and architecture-independent work-tile transport.""" from abc import ABC, abstractmethod @@ -15,6 +12,7 @@ from ...helpers.device_workspace import DeviceWorkspace from ...helpers.smem_workspace import SmemWorkspace + WorkIdAcquisitionMode = Literal["grid_stride", "atomic_counter", "cluster_launch_control"] @@ -55,9 +53,7 @@ def __init__( self._pipeline = scheduler_pipeline self._smem_buffer = smem_buffer self._work_tile_type = work_tile_type - self._consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, num_stages - ) + self._consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, num_stages) def __extract_mlir_values__(self) -> list: return extract_mlir_values(self._consumer_state) @@ -122,9 +118,7 @@ def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: ) work_tile_type = self.work_tile_type work_tile_field_count = work_tile_type.storage_field_count - smem_workspace.register_mbarrier( - self.pipeline_mbarriers_region, self.num_scheduler_stages * 2 - ) + smem_workspace.register_mbarrier(self.pipeline_mbarriers_region, self.num_scheduler_stages * 2) smem_workspace.register_tensor( self.work_tiles_region, work_tile_type.storage_dtype, @@ -138,16 +132,12 @@ def register_device_workspace(self, device_workspace: DeviceWorkspace) -> None: pass @cute.jit - def create_scheduler_pipelines( - self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer - ) -> None: + def create_scheduler_pipelines(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: """Create CTA-lifetime scheduler pipelines and transport state.""" self._pipeline = pipeline.PipelineAsync.create( num_stages=self.num_scheduler_stages, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, self.num_scheduler_consumer_threads - ), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, self.num_scheduler_consumer_threads), barrier_storage=smem_workspace.ptr(self.pipeline_mbarriers_region, smem_base), defer_sync=True, ) @@ -168,13 +158,9 @@ def make_consumer(self) -> SchedulerConsumer: @cute.jit def publish_work(self, work_tile: SchedulerWorkTileBase) -> None: """Publish one work tile through the common transport pipeline.""" - copy_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), work_tile.storage_dtype, num_bits_per_copy=128 - ) + copy_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), work_tile.storage_dtype, num_bits_per_copy=128) self._pipeline.producer_acquire(self._producer_state) - cute.copy( - copy_atom, work_tile.to_rmem(), self._smem_buffer[(None, self._producer_state.index)] - ) + cute.copy(copy_atom, work_tile.to_rmem(), self._smem_buffer[(None, self._producer_state.index)]) cute.arch.fence_proxy("async.shared", space="cta") self._pipeline.producer_commit(self._producer_state) self._producer_state.advance() @@ -202,9 +188,7 @@ def __new_from_mlir_values__(self, values: list) -> "SchedulerBase": return result @abstractmethod - def get_grid_shape( - self, *, max_active_clusters: Optional[int] = None, problem_desc: Any = None - ): + def get_grid_shape(self, *, max_active_clusters: Optional[int] = None, problem_desc: Any = None): """Return the launch grid from static scheduler policy.""" ... diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/clc_mixed_cga.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/clc_mixed_cga.py new file mode 100644 index 000000000000..896f1d15989d --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/clc_mixed_cga.py @@ -0,0 +1,162 @@ +"""Hardware CLC work-ID acquisition for preferred/fallback (mixed-CGA) cluster launches.""" + +from typing import List, Optional + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +from cutlass._mlir import ir +from cutlass.cutlass_dsl import Boolean, Int32, extract_mlir_values, new_from_mlir_values + +from ...helpers.device_workspace import DeviceWorkspace +from ...helpers.smem_workspace import SmemWorkspace +from .non_clc_mixed_cga import NonClcMixedCgaConfig +from .work_id_claim import ClcMixedCgaWorkIdState, claim_work_id + + +class ClcMixedCgaSchedulerWorker: + """Claim linear work IDs from hardware cluster launch control on a mixed-CGA launch. + + Same surface as ``NonClcMixedCgaSchedulerWorker`` plus ``claimed_is_valid``: a failed + ``try_cancel`` means the grid is exhausted and the caller must emit the DONE tile. + """ + + cluster_pipeline_mbarriers_region = "clc_mixed_cga.cluster_pipeline_mbarriers" + clc_response_region = "clc_mixed_cga.clc_response" + + def __init__(self, *, config: NonClcMixedCgaConfig, stream_count: int = 1) -> None: + if stream_count != 1: + raise ValueError("Cluster launch control exposes exactly one hardware work-ID stream.") + self.config = config + self.work_id_mode = "cluster_launch_control" + self.stream_count = stream_count + + def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: + """Register the cluster claim barriers and the 16-byte CLC response slot.""" + smem_workspace.register_mbarrier(self.cluster_pipeline_mbarriers_region, 2) + smem_workspace.register_tensor(self.clc_response_region, cutlass.Int32, (4,), byte_alignment=16) + + def register_device_workspace(self, device_workspace: DeviceWorkspace) -> None: + """CLC keeps no GMEM scheduler state.""" + pass + + @cute.jit + def assign_device_members( + self, + *, + is_fallback_cluster: Optional[Boolean], + block_idx, + smem_workspace: SmemWorkspace, + smem_base: cute.Pointer, + device_workspace: DeviceWorkspace, + ) -> None: + """Bind this CTA's active cluster to the CLC claim pipeline and seed the bootstrap ticket.""" + if cutlass.const_expr(self.config.is_mixed and is_fallback_cluster is None): + raise ValueError("is_fallback_cluster is required for a true mixed-CGA launch.") + + preferred_m, preferred_n = self.config.preferred_cluster_shape + active_cluster_m = preferred_m + active_cluster_n = preferred_n + if cutlass.const_expr(self.config.is_mixed): + if is_fallback_cluster: + active_cluster_m = Int32(self.config.fallback_cluster_shape[0]) + active_cluster_n = Int32(self.config.fallback_cluster_shape[1]) + + cta_coord_in_active_cluster = ( + Int32(block_idx[0]) % Int32(active_cluster_m), + Int32(block_idx[1]) % Int32(active_cluster_n), + Int32(0), + ) + active_cluster_size = active_cluster_m * active_cluster_n + cluster_pipeline = pipeline.PipelineAsync.create( + num_stages=1, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * active_cluster_size), + barrier_storage=smem_workspace.ptr(self.cluster_pipeline_mbarriers_region, smem_base), + defer_sync=True, + ) + self._work_id_state = ClcMixedCgaWorkIdState( + response_pending=Boolean(True), + grid_m=Int32(block_idx[0]), + grid_n=Int32(block_idx[1]), + grid_l=Int32(block_idx[2]), + response_is_valid=Boolean(True), + cta_coord_in_cluster=cta_coord_in_active_cluster, + cluster_pipeline=cluster_pipeline, + producer_state=pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1), + consumer_state=pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1), + is_leader_cta=(cta_coord_in_active_cluster[0] + cta_coord_in_active_cluster[1]) == Int32(0), + response_pointer=smem_workspace.ptr(self.clc_response_region, smem_base), + cluster_size=active_cluster_size, + ) + self.cta_coord_in_preferred_cluster = ( + Int32(block_idx[0]) % Int32(preferred_m), + Int32(block_idx[1]) % Int32(preferred_n), + Int32(0), + ) + self.claimed_stream_index = Int32(0) + self.claimed_is_valid = Boolean(True) + + @cute.jit + def initialize_fallback_group(self) -> None: + """No fallback registration: hardware assigns fallback clusters to slots.""" + pass + + @cute.jit + def claim_next_work(self, stream_index=0) -> Int32: + """Claim the next ticket; publish its validity and this CTA's preferred-cluster coordinate.""" + work_id, self._work_id_state = claim_work_id(self._work_id_state) + preferred_m, preferred_n = self.config.preferred_cluster_shape + self.cta_coord_in_preferred_cluster = ( + work_id.grid_m % Int32(preferred_m), + work_id.grid_n % Int32(preferred_n), + Int32(0), + ) + self.claimed_is_valid = work_id.is_valid + self.claimed_stream_index = Int32(0) + # Exhausted grid: hand the mapper a harmless in-range ID; the caller substitutes the DONE tile. + linear_work_id = Int32(0) + if work_id.is_valid: + linear_work_id = work_id.grid_l + return linear_work_id + + @cute.jit + def produce_tail(self) -> None: + """Leader waits for the cluster's final release, since peers release on rank 0's empty barrier.""" + work_id_state = self._work_id_state + if self._work_id_state.is_leader_cta: + self._work_id_state.cluster_pipeline.producer_tail(self._work_id_state.producer_state) + else: + self._work_id_state = work_id_state + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + values.extend(extract_mlir_values(self._work_id_state)) + values.extend(extract_mlir_values(self.cta_coord_in_preferred_cluster)) + values.extend(extract_mlir_values(self.claimed_stream_index)) + values.extend(extract_mlir_values(self.claimed_is_valid)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "ClcMixedCgaSchedulerWorker": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) + value_index += field_value_count + return result + + result = type(self)(config=self.config, stream_count=self.stream_count) + result._work_id_state = rebuild(self._work_id_state) + result.cta_coord_in_preferred_cluster = rebuild(self.cta_coord_in_preferred_cluster) + result.claimed_stream_index = rebuild(self.claimed_stream_index) + result.claimed_is_valid = rebuild(self.claimed_is_valid) + if value_index != len(values): + raise ValueError( + f"ClcMixedCgaSchedulerWorker MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return result + + +__all__ = ["ClcMixedCgaSchedulerWorker"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/fc12_mapping.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/fc12_mapping.py index 14d186d45fe9..0bdd9b9ce720 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/fc12_mapping.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/fc12_mapping.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """FC12 work-tile ABI and grouped or phase-interleaved task mapping.""" import dataclasses @@ -16,11 +13,16 @@ from ...helpers.utils import padded_expert_rows from .base import SchedulerWorkTileBase + phase_bits = 16 phase_mask = (1 << phase_bits) - 1 peek_ready_bit = 1 << phase_bits + + + + class Fc12WorkTileState(IntEnum): """Sentinel values carried in the expert index field.""" @@ -269,9 +271,7 @@ def __new_from_mlir_values__(self, values: List[ir.Value]) -> "_Fc12TaskCursorSt rebuilt = [] for field in self._fields(): field_value_count = len(extract_mlir_values(field)) - rebuilt.append( - new_from_mlir_values(field, values[value_index : value_index + field_value_count]) - ) + rebuilt.append(new_from_mlir_values(field, values[value_index : value_index + field_value_count])) value_index += field_value_count if value_index != len(values): raise ValueError( @@ -341,11 +341,7 @@ def __extract_mlir_values__(self) -> List[ir.Value]: values: List[ir.Value] = [] if isinstance(self.expert_count, Int32): values.extend(extract_mlir_values(self.expert_count)) - token_counts = ( - self.expert_token_sizes - if self.expert_token_sizes is not None - else self.expert_token_prefix_sum - ) + token_counts = self.expert_token_sizes if self.expert_token_sizes is not None else self.expert_token_prefix_sum values.extend(extract_mlir_values(token_counts)) values.extend(extract_mlir_values(self.cursor_state)) if isinstance(self.num_fc1_intermediate_blocks, Int32): @@ -360,17 +356,11 @@ def __new_from_mlir_values__(self, values: List[ir.Value]) -> "Fc12TaskMappingSt def rebuild(field): nonlocal value_index field_value_count = len(extract_mlir_values(field)) - result = new_from_mlir_values( - field, values[value_index : value_index + field_value_count] - ) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) value_index += field_value_count return result - expert_count = ( - rebuild(self.expert_count) - if isinstance(self.expert_count, Int32) - else self.expert_count - ) + expert_count = rebuild(self.expert_count) if isinstance(self.expert_count, Int32) else self.expert_count if self.expert_token_sizes is not None: expert_token_sizes = rebuild(self.expert_token_sizes) expert_token_prefix_sum = None @@ -422,7 +412,29 @@ def create_fc12_task_mapping_state( expert_token_prefix_sum: Optional[cute.Tensor], ) -> Fc12TaskMappingState: """Create the register-resident state for one CTA's FC12 mapper.""" - cursor_state = _Fc12TaskCursorState( + cursor_state = _make_initial_fc12_cursor() + mapping_cluster_tile_n = mapping_cluster_shape_mn[1] * mapping_cta_tile_shape_mnk[1] + num_fc1_intermediate_blocks = (intermediate_gateup_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n + num_fc2_hidden_blocks = (hidden_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n + return Fc12TaskMappingState( + expert_count=expert_count, + mapping_cta_tile_shape_mnk=mapping_cta_tile_shape_mnk, + mapping_cluster_shape_mn=mapping_cluster_shape_mn, + group_hint=group_hint, + token_padding_block=token_padding_block, + sf_padding_block=sf_padding_block, + is_swap_ab=is_swap_ab, + expert_token_sizes=expert_token_sizes, + expert_token_prefix_sum=expert_token_prefix_sum, + cursor_state=cursor_state, + num_fc1_intermediate_blocks=num_fc1_intermediate_blocks, + num_fc2_hidden_blocks=num_fc2_hidden_blocks, + ) + + +def _make_initial_fc12_cursor() -> _Fc12TaskCursorState: + """Cursor positioned before the first group; the mapper seeks forward from here.""" + return _Fc12TaskCursorState( current_group_first_expert=Int32(0), current_group_last_expert_exclusive=Int32(0), current_phase=Int32(BlockPhase.Linear1), @@ -442,25 +454,13 @@ def create_fc12_task_mapping_state( current_token_block_count=Int32(0), current_expert_token_count=Int32(0), ) - mapping_cluster_tile_n = mapping_cluster_shape_mn[1] * mapping_cta_tile_shape_mnk[1] - num_fc1_intermediate_blocks = ( - intermediate_gateup_size + mapping_cluster_tile_n - 1 - ) // mapping_cluster_tile_n - num_fc2_hidden_blocks = (hidden_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n - return Fc12TaskMappingState( - expert_count=expert_count, - mapping_cta_tile_shape_mnk=mapping_cta_tile_shape_mnk, - mapping_cluster_shape_mn=mapping_cluster_shape_mn, - group_hint=group_hint, - token_padding_block=token_padding_block, - sf_padding_block=sf_padding_block, - is_swap_ab=is_swap_ab, - expert_token_sizes=expert_token_sizes, - expert_token_prefix_sum=expert_token_prefix_sum, - cursor_state=cursor_state, - num_fc1_intermediate_blocks=num_fc1_intermediate_blocks, - num_fc2_hidden_blocks=num_fc2_hidden_blocks, - ) + + +@cute.jit +def rewind_fc12_mapping_cursor(mapping_state: Fc12TaskMappingState) -> Fc12TaskMappingState: + """Restart the forward-only cursor so a lower linear ID (a non-monotonic CLC ticket) can be mapped.""" + mapping_state.cursor_state = _make_initial_fc12_cursor() + return mapping_state @cute.jit @@ -486,11 +486,7 @@ def _first_matching_lane(predicate) -> Int32: @cute.jit def _load_expert_batch_metrics( - mapping_state: Fc12TaskMappingState, - batch_base: Int32, - active_begin: Int32, - active_end: Int32, - lane_idx: Int32, + mapping_state: Fc12TaskMappingState, batch_base: Int32, active_begin: Int32, active_end: Int32, lane_idx: Int32 ) -> Tuple[Int32, Int32, Int32, Int32, Int32, Int32, Int32]: expert_idx = batch_base + lane_idx token_count = Int32(0) @@ -582,10 +578,7 @@ def _sum_expert_range( @cute.jit def _build_group_range( - mapping_state: Fc12TaskMappingState, - group_first_expert: Int32, - base_fc1_tiles: Int32, - base_fc2_tiles: Int32, + mapping_state: Fc12TaskMappingState, group_first_expert: Int32, base_fc1_tiles: Int32, base_fc2_tiles: Int32 ) -> Tuple[Int32, Int32, Int32]: lane_idx = Int32(cute.arch.lane_idx()) group_threshold = base_fc1_tiles + Int32(mapping_state.group_hint) @@ -596,11 +589,7 @@ def _build_group_range( while batch_base < mapping_state.expert_count and cumulative_fc1_tiles < group_threshold: (lane_expert_idx, _, _, _, _, lane_fc1_tiles, lane_fc2_tiles) = _load_expert_batch_metrics( - mapping_state, - batch_base, - group_first_expert, - Int32(mapping_state.expert_count), - lane_idx, + mapping_state, batch_base, group_first_expert, Int32(mapping_state.expert_count), lane_idx ) fc1_prefix = _warp_inclusive_sum(lane_fc1_tiles, lane_idx) reaches_threshold = ( @@ -613,14 +602,10 @@ def _build_group_range( included_fc2_tiles = Int32(0) fc2_batch_tiles = Int32(cute.arch.warp_redux_sync(included_fc2_tiles, "add")) if selected_lane >= Int32(0): - cumulative_fc1_tiles = cumulative_fc1_tiles + Int32( - cute.arch.shuffle_sync(fc1_prefix, selected_lane) - ) + cumulative_fc1_tiles = cumulative_fc1_tiles + Int32(cute.arch.shuffle_sync(fc1_prefix, selected_lane)) group_last_expert = batch_base + selected_lane + Int32(1) else: - cumulative_fc1_tiles = cumulative_fc1_tiles + Int32( - cute.arch.shuffle_sync(fc1_prefix, Int32(31)) - ) + cumulative_fc1_tiles = cumulative_fc1_tiles + Int32(cute.arch.shuffle_sync(fc1_prefix, Int32(31))) batch_base = batch_base + Int32(32) group_last_expert = cutlass.min(batch_base, Int32(mapping_state.expert_count)) cumulative_fc2_tiles = cumulative_fc2_tiles + fc2_batch_tiles @@ -644,9 +629,7 @@ def _advance_group(mapping_state: Fc12TaskMappingState) -> _Fc12TaskCursorState: iket.range_pop() cursor.current_data_cumulative = cursor.current_data_cumulative + residual_data_rows cursor.current_sf_cumulative = cursor.current_sf_cumulative + residual_sf_rows - cursor.current_token_block_cumulative = ( - cursor.current_token_block_cumulative + residual_token_blocks - ) + cursor.current_token_block_cumulative = cursor.current_token_block_cumulative + residual_token_blocks cursor.group_start_data_cumulative = cursor.current_data_cumulative cursor.group_start_sf_cumulative = cursor.current_sf_cumulative @@ -657,19 +640,15 @@ def _advance_group(mapping_state: Fc12TaskMappingState) -> _Fc12TaskCursorState: cursor.current_group_first_expert = cursor.current_group_last_expert_exclusive iket.range_push("scheduler.group_scan") - (cursor.current_group_last_expert_exclusive, cumulative_fc1_tiles, cumulative_fc2_tiles) = ( - _build_group_range( - mapping_state, cursor.current_group_first_expert, base_fc1_tiles, base_fc2_tiles - ) + (cursor.current_group_last_expert_exclusive, cumulative_fc1_tiles, cumulative_fc2_tiles) = _build_group_range( + mapping_state, cursor.current_group_first_expert, base_fc1_tiles, base_fc2_tiles ) iket.range_pop() cursor.cumulative_fc1_tiles_at_group_end = cumulative_fc1_tiles cursor.cumulative_fc2_tiles_at_group_end = cumulative_fc2_tiles group_start_tile = cursor.current_group_end cursor.current_group_fc1_subphase_end = group_start_tile + cumulative_fc1_tiles - base_fc1_tiles - cursor.current_group_end = ( - cursor.current_group_fc1_subphase_end + cumulative_fc2_tiles - base_fc2_tiles - ) + cursor.current_group_end = cursor.current_group_fc1_subphase_end + cumulative_fc2_tiles - base_fc2_tiles cursor.current_phase = Int32(BlockPhase.Linear1) cursor.current_expert_idx = cursor.current_group_first_expert - Int32(1) @@ -682,9 +661,7 @@ def _advance_group(mapping_state: Fc12TaskMappingState) -> _Fc12TaskCursorState: @cute.jit -def _seek_expert_for_work_id( - linear_work_id: Int32, mapping_state: Fc12TaskMappingState -) -> _Fc12TaskCursorState: +def _seek_expert_for_work_id(linear_work_id: Int32, mapping_state: Fc12TaskMappingState) -> _Fc12TaskCursorState: cursor = mapping_state.cursor_state base_tile_end = cursor.current_expert_tile_end @@ -701,9 +678,7 @@ def _seek_expert_for_work_id( ) base_token_block_cumulative = base_token_block_cumulative + cursor.current_token_block_count - search_begin = cutlass.max( - cursor.current_expert_idx + Int32(1), cursor.current_group_first_expert - ) + search_begin = cutlass.max(cursor.current_expert_idx + Int32(1), cursor.current_group_first_expert) batch_base = (search_begin // Int32(32)) * Int32(32) selected_expert = Int32(-1) selected_token_count = Int32(0) @@ -726,11 +701,7 @@ def _seek_expert_for_work_id( lane_fc1_tiles, lane_fc2_tiles, ) = _load_expert_batch_metrics( - mapping_state, - batch_base, - search_begin, - cursor.current_group_last_expert_exclusive, - lane_idx, + mapping_state, batch_base, search_begin, cursor.current_group_last_expert_exclusive, lane_idx ) lane_phase_tiles = lane_fc1_tiles if cursor.current_phase == Int32(BlockPhase.Linear2): @@ -760,17 +731,13 @@ def _seek_expert_for_work_id( selected_token_count = Int32(cute.arch.shuffle_sync(lane_token_count, selected_lane)) selected_token_blocks = Int32(cute.arch.shuffle_sync(lane_token_blocks, selected_lane)) selected_phase_tiles = Int32(cute.arch.shuffle_sync(lane_phase_tiles, selected_lane)) - selected_tile_end = base_tile_end + Int32( - cute.arch.shuffle_sync(tile_prefix, selected_lane) - ) + selected_tile_end = base_tile_end + Int32(cute.arch.shuffle_sync(tile_prefix, selected_lane)) selected_tile_start = selected_tile_end - selected_phase_tiles selected_data_rows = Int32(cute.arch.shuffle_sync(lane_data_rows, selected_lane)) selected_sf_rows = Int32(cute.arch.shuffle_sync(lane_sf_rows, selected_lane)) selected_data_cumulative = base_data_cumulative + batch_data_rows - selected_data_rows selected_sf_cumulative = base_sf_cumulative + batch_sf_rows - selected_sf_rows - selected_token_block_cumulative = ( - base_token_block_cumulative + batch_token_blocks - selected_token_blocks - ) + selected_token_block_cumulative = base_token_block_cumulative + batch_token_blocks - selected_token_blocks else: base_tile_end = base_tile_end + Int32(cute.arch.shuffle_sync(tile_prefix, Int32(31))) base_data_cumulative = base_data_cumulative + batch_data_rows @@ -793,9 +760,7 @@ def _seek_expert_for_work_id( @cute.jit def _decode_inside_expert( - linear_work_id: Int32, - cta_id_in_mapping_cluster: cute.Coord, - mapping_state: Fc12TaskMappingState, + linear_work_id: Int32, cta_id_in_mapping_cluster: cute.Coord, mapping_state: Fc12TaskMappingState ) -> SchedulerWorkTileBase: cursor = mapping_state.cursor_state cta_tile_m = mapping_state.mapping_cta_tile_shape_mnk[0] @@ -805,22 +770,16 @@ def _decode_inside_expert( cluster_output_block_idx = Int32(0) if cursor.current_phase == Int32(BlockPhase.Linear1): cluster_token_block_idx = local_work_id // mapping_state.num_fc1_intermediate_blocks - cluster_output_block_idx = ( - local_work_id - cluster_token_block_idx * mapping_state.num_fc1_intermediate_blocks - ) + cluster_output_block_idx = local_work_id - cluster_token_block_idx * mapping_state.num_fc1_intermediate_blocks else: cluster_token_block_idx = local_work_id // mapping_state.num_fc2_hidden_blocks - cluster_output_block_idx = ( - local_work_id - cluster_token_block_idx * mapping_state.num_fc2_hidden_blocks - ) + cluster_output_block_idx = local_work_id - cluster_token_block_idx * mapping_state.num_fc2_hidden_blocks cta_token_block_idx = ( - cluster_token_block_idx * mapping_state.mapping_cluster_shape_mn[0] - + cta_id_in_mapping_cluster[0] + cluster_token_block_idx * mapping_state.mapping_cluster_shape_mn[0] + cta_id_in_mapping_cluster[0] ) cta_output_block_idx = ( - cluster_output_block_idx * mapping_state.mapping_cluster_shape_mn[1] - + cta_id_in_mapping_cluster[1] + cluster_output_block_idx * mapping_state.mapping_cluster_shape_mn[1] + cta_id_in_mapping_cluster[1] ) token_start = cta_token_block_idx * Int32(cta_tile_m) remaining_tokens = cutlass.max(cursor.current_expert_token_count - token_start, Int32(0)) @@ -840,13 +799,9 @@ def _decode_inside_expert( cluster_tile_m = mapping_state.mapping_cluster_shape_mn[0] * cta_tile_m cluster_token_start = cluster_token_block_idx * Int32(cluster_tile_m) - remaining_cluster_tokens = cutlass.max( - cursor.current_expert_token_count - cluster_token_start, Int32(0) - ) + remaining_cluster_tokens = cutlass.max(cursor.current_expert_token_count - cluster_token_start, Int32(0)) valid_tokens_in_cluster_tile = cutlass.min(remaining_cluster_tokens, Int32(cluster_tile_m)) - valid_tokens_in_cta_cluster_tile = ( - valid_tokens_in_cta_tile << Int32(16) - ) | valid_tokens_in_cluster_tile + valid_tokens_in_cta_cluster_tile = (valid_tokens_in_cta_tile << Int32(16)) | valid_tokens_in_cluster_tile return NonSwapAbFc12WorkTileInfo( expert_idx=cursor.current_expert_idx, tile_m_idx=cta_token_block_idx, @@ -861,9 +816,7 @@ def _decode_inside_expert( @cute.jit def map_fc12_linear_work_id( - linear_work_id: Int32, - cta_id_in_mapping_cluster: cute.Coord, - mapping_state: Fc12TaskMappingState, + linear_work_id: Int32, cta_id_in_mapping_cluster: cute.Coord, mapping_state: Fc12TaskMappingState ) -> Tuple[SchedulerWorkTileBase, Fc12TaskMappingState]: """Map one monotonically increasing scalar ID to an FC12 work tile.""" cursor = mapping_state.cursor_state @@ -939,9 +892,7 @@ def _runtime_fields(self) -> Tuple: @cute.jit def clone(self) -> "_PhaseFc12CursorState": """Fork this cursor while sharing its immutable initial SSA values.""" - return type(self)( - *self._runtime_fields(), blocks_per_token_block=self.blocks_per_token_block - ) + return type(self)(*self._runtime_fields(), blocks_per_token_block=self.blocks_per_token_block) def __extract_mlir_values__(self) -> List[ir.Value]: values: List[ir.Value] = [] @@ -954,9 +905,7 @@ def __new_from_mlir_values__(self, values: List[ir.Value]) -> "_PhaseFc12CursorS rebuilt_fields = [] for field in self._runtime_fields(): field_value_count = len(extract_mlir_values(field)) - rebuilt_fields.append( - new_from_mlir_values(field, values[value_index : value_index + field_value_count]) - ) + rebuilt_fields.append(new_from_mlir_values(field, values[value_index : value_index + field_value_count])) value_index += field_value_count if value_index != len(values): raise ValueError( @@ -982,7 +931,11 @@ def __init__( fc2_cursor: _PhaseFc12CursorState, num_fc1_intermediate_blocks: int, num_fc2_hidden_blocks: int, + dfc1_m_group: int = 1, ) -> None: + if type(dfc1_m_group) is not int or dfc1_m_group not in (1, 16): + raise ValueError("dfc1_m_group must be 1 or 16") + self.dfc1_m_group = dfc1_m_group self.expert_count = expert_count self.mapping_cta_tile_shape_mnk = mapping_cta_tile_shape_mnk self.mapping_cluster_shape_mn = mapping_cluster_shape_mn @@ -995,6 +948,12 @@ def __init__( self.fc2_cursor = fc2_cursor self.num_fc1_intermediate_blocks = num_fc1_intermediate_blocks self.num_fc2_hidden_blocks = num_fc2_hidden_blocks + self.fc2_claims_per_token_block = num_fc2_hidden_blocks + if fc2_cursor.blocks_per_token_block != self.fc2_claims_per_token_block: + raise ValueError( + "FC2 cursor blocks_per_token_block does not match the selected " + "work-ID domain." + ) @property def mapping_cluster_tile_m(self) -> int: @@ -1002,27 +961,19 @@ def mapping_cluster_tile_m(self) -> int: def __extract_mlir_values__(self) -> List[ir.Value]: values: List[ir.Value] = [] - token_counts = ( - self.expert_token_sizes - if self.expert_token_sizes is not None - else self.expert_token_prefix_sum - ) + token_counts = self.expert_token_sizes if self.expert_token_sizes is not None else self.expert_token_prefix_sum values.extend(extract_mlir_values(token_counts)) for field in (self.fc1_cursor, self.fc2_cursor): values.extend(extract_mlir_values(field)) return values - def __new_from_mlir_values__( - self, values: List[ir.Value] - ) -> "PhaseInterleavedFc12MappingState": + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "PhaseInterleavedFc12MappingState": value_index = 0 def rebuild(field): nonlocal value_index field_value_count = len(extract_mlir_values(field)) - result = new_from_mlir_values( - field, values[value_index : value_index + field_value_count] - ) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) value_index += field_value_count return result @@ -1045,6 +996,7 @@ def rebuild(field): fc2_cursor=rebuild(self.fc2_cursor), num_fc1_intermediate_blocks=self.num_fc1_intermediate_blocks, num_fc2_hidden_blocks=self.num_fc2_hidden_blocks, + dfc1_m_group=self.dfc1_m_group, ) if value_index != len(values): raise ValueError( @@ -1081,13 +1033,13 @@ def create_phase_interleaved_fc12_mapping_state( is_swap_ab: bool, expert_token_sizes: Optional[cute.Tensor], expert_token_prefix_sum: Optional[cute.Tensor], + dfc1_m_group: int = 1, ) -> PhaseInterleavedFc12MappingState: """Create independent monotonic mapping cursors for the FC1 and FC2 streams.""" mapping_cluster_tile_n = mapping_cluster_shape_mn[1] * mapping_cta_tile_shape_mnk[1] - num_fc1_intermediate_blocks = ( - intermediate_gateup_size + mapping_cluster_tile_n - 1 - ) // mapping_cluster_tile_n + num_fc1_intermediate_blocks = (intermediate_gateup_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n num_fc2_hidden_blocks = (hidden_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n + fc2_claims_per_token_block = num_fc2_hidden_blocks return PhaseInterleavedFc12MappingState( expert_count=expert_count, mapping_cta_tile_shape_mnk=mapping_cta_tile_shape_mnk, @@ -1098,9 +1050,10 @@ def create_phase_interleaved_fc12_mapping_state( expert_token_sizes=expert_token_sizes, expert_token_prefix_sum=expert_token_prefix_sum, fc1_cursor=_make_phase_cursor(num_fc1_intermediate_blocks), - fc2_cursor=_make_phase_cursor(num_fc2_hidden_blocks), + fc2_cursor=_make_phase_cursor(fc2_claims_per_token_block), num_fc1_intermediate_blocks=num_fc1_intermediate_blocks, num_fc2_hidden_blocks=num_fc2_hidden_blocks, + dfc1_m_group=dfc1_m_group, ) @@ -1129,9 +1082,9 @@ def _advance_phase_cursor( token_count = prefix_end - prefix_begin cursor.current_expert_token_count = token_count - cursor.current_token_block_count = ( - token_count + Int32(mapping_state.mapping_cluster_tile_m - 1) - ) // Int32(mapping_state.mapping_cluster_tile_m) + cursor.current_token_block_count = (token_count + Int32(mapping_state.mapping_cluster_tile_m - 1)) // Int32( + mapping_state.mapping_cluster_tile_m + ) cursor.expert_tile_start = cursor.expert_tile_end cursor.expert_tile_end = cursor.expert_tile_start + cursor.current_token_block_count * Int32( cursor.blocks_per_token_block @@ -1141,9 +1094,7 @@ def _advance_phase_cursor( @cute.jit def _seek_phase_cursor( - linear_work_id: Int32, - cursor: _PhaseFc12CursorState, - mapping_state: PhaseInterleavedFc12MappingState, + linear_work_id: Int32, cursor: _PhaseFc12CursorState, mapping_state: PhaseInterleavedFc12MappingState ) -> _PhaseFc12CursorState: expert_tile_end = cursor.expert_tile_end next_expert_idx = cursor.expert_idx + Int32(1) @@ -1160,9 +1111,7 @@ def resolve_phase_interleaved_fc1_claim_target( ) -> Tuple[Int32, Boolean]: """Resolve the static FC1 watermark or the smaller runtime stream extent.""" target_work_id = minimum_claim_count - Int32(1) - probe_cursor = _seek_phase_cursor( - target_work_id, mapping_state.fc1_cursor.clone(), mapping_state - ) + probe_cursor = _seek_phase_cursor(target_work_id, mapping_state.fc1_cursor.clone(), mapping_state) stream_ends_before_target = target_work_id >= probe_cursor.expert_tile_end claim_target = minimum_claim_count if stream_ends_before_target: @@ -1179,17 +1128,37 @@ def _decode_phase_work_id( mapping_state: PhaseInterleavedFc12MappingState, ) -> SchedulerWorkTileBase: local_work_id = linear_work_id - cursor.expert_tile_start - cluster_token_block_idx = local_work_id // Int32(cursor.blocks_per_token_block) + cluster_token_block_idx = local_work_id // Int32( + cursor.blocks_per_token_block + ) cluster_output_block_idx = local_work_id - cluster_token_block_idx * Int32( cursor.blocks_per_token_block ) + if cutlass.const_expr(mapping_state.dfc1_m_group > 1): + if phase == Int32(BlockPhase.Linear2): + # Within each small M strip, visit its rows for a fixed N weight + # tile before advancing N. Only the per-expert rectangle is + # permuted: expert ranges, real row/ready indices, and task counts + # remain unchanged. A short final strip uses its actual row count + # and therefore adds neither padding work nor duplicate tiles. + strip_span = Int32(mapping_state.dfc1_m_group * cursor.blocks_per_token_block) + strip = local_work_id // strip_span + strip_m_begin = strip * Int32(mapping_state.dfc1_m_group) + strip_rows = cutlass.max(Int32(1), cutlass.min( + Int32(mapping_state.dfc1_m_group), + cursor.current_token_block_count - strip_m_begin, + )) + in_strip = local_work_id - strip * strip_span + cluster_output_block_idx = in_strip // strip_rows + cluster_token_block_idx = ( + strip_m_begin + in_strip - cluster_output_block_idx * strip_rows + ) + cta_token_block_idx = ( - cluster_token_block_idx * Int32(mapping_state.mapping_cluster_shape_mn[0]) - + cta_id_in_mapping_cluster[0] + cluster_token_block_idx * Int32(mapping_state.mapping_cluster_shape_mn[0]) + cta_id_in_mapping_cluster[0] ) cta_output_block_idx = ( - cluster_output_block_idx * Int32(mapping_state.mapping_cluster_shape_mn[1]) - + cta_id_in_mapping_cluster[1] + cluster_output_block_idx * Int32(mapping_state.mapping_cluster_shape_mn[1]) + cta_id_in_mapping_cluster[1] ) cta_tile_m = mapping_state.mapping_cta_tile_shape_mnk[0] @@ -1211,9 +1180,7 @@ def _decode_phase_work_id( cluster_tile_m = mapping_state.mapping_cluster_shape_mn[0] * cta_tile_m cluster_token_start = cluster_token_block_idx * Int32(cluster_tile_m) - remaining_cluster_tokens = cutlass.max( - cursor.current_expert_token_count - cluster_token_start, Int32(0) - ) + remaining_cluster_tokens = cutlass.max(cursor.current_expert_token_count - cluster_token_start, Int32(0)) valid_tokens_in_cluster_tile = cutlass.min(remaining_cluster_tokens, Int32(cluster_tile_m)) return NonSwapAbFc12WorkTileInfo( expert_idx=cursor.expert_idx, @@ -1222,8 +1189,7 @@ def _decode_phase_work_id( cumulative_data_physical_row=cursor.data_cumulative, cumulative_sf_physical_row=cursor.sf_cumulative, cumulative_token_block_count=cursor.token_block_cumulative, - valid_tokens_in_cta_cluster_tile=(valid_tokens_in_cta_tile << Int32(16)) - | valid_tokens_in_cluster_tile, + valid_tokens_in_cta_cluster_tile=(valid_tokens_in_cta_tile << Int32(16)) | valid_tokens_in_cluster_tile, phase_and_flags=phase, ) @@ -1261,6 +1227,8 @@ def map_phase_interleaved_fc12_work_id( return work_tile, stream_has_work, mapping_state + + __all__ = [ "BlockPhase", "Fc12TaskMappingState", @@ -1275,4 +1243,5 @@ def map_phase_interleaved_fc12_work_id( "map_phase_interleaved_fc12_work_id", "peek_ready_bit", "resolve_phase_interleaved_fc1_claim_target", + "rewind_fc12_mapping_cursor", ] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/fc12_scheduler.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/fc12_scheduler.py index 75ef3eccaabd..9fd349cd8048 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/fc12_scheduler.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/fc12_scheduler.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Composable grouped and phase-interleaved FC12 schedulers.""" import math @@ -16,6 +13,7 @@ from ...helpers.smem_workspace import SmemWorkspace from ...helpers.utils import ceil_div from .base import SchedulerBase, SchedulerWorkTileBase, WorkIdAcquisitionMode +from .clc_mixed_cga import ClcMixedCgaSchedulerWorker from .fc12_mapping import ( BlockPhase, NonSwapAbFc12WorkTileInfo, @@ -26,6 +24,7 @@ map_fc12_linear_work_id, map_phase_interleaved_fc12_work_id, resolve_phase_interleaved_fc1_claim_target, + rewind_fc12_mapping_cursor, ) from .non_clc_mixed_cga import NonClcMixedCgaConfig, NonClcMixedCgaSchedulerWorker @@ -70,16 +69,10 @@ def minimum_phase_interleave_hint( return max(1, ceil_div(required_fc1_claims, launch_cluster_cnt_merge_as_preferred)) -def _to_fc12_mapping_cta_coord( - cta_coord_in_preferred_cluster: cute.Coord, is_swap_ab: bool -) -> cute.Coord: +def _to_fc12_mapping_cta_coord(cta_coord_in_preferred_cluster: cute.Coord, is_swap_ab: bool) -> cute.Coord: if cutlass.const_expr(not is_swap_ab): return cta_coord_in_preferred_cluster - return ( - cta_coord_in_preferred_cluster[1], - cta_coord_in_preferred_cluster[0], - cta_coord_in_preferred_cluster[2], - ) + return (cta_coord_in_preferred_cluster[1], cta_coord_in_preferred_cluster[0], cta_coord_in_preferred_cluster[2]) class BlackwellFusedFc12Scheduler(SchedulerBase): @@ -109,6 +102,7 @@ def impl_desc_require(cls) -> dict[str, type]: "work_id_mode": str, "is_swap_ab": bool, **_mixed_cga_impl_requirements(), + "compute_cluster_count": OptionalRequirement(int), } def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: @@ -127,14 +121,11 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.work_id_mode: WorkIdAcquisitionMode = impl_desc["work_id_mode"] self.is_swap_ab = impl_desc["is_swap_ab"] self.non_clc_mixed_cga_config = _make_non_clc_mixed_cga_config(impl_desc) - self.launch_cluster_cnt_merge_as_preferred = ( - self.non_clc_mixed_cga_config.launch_cluster_cnt_merge_as_preferred - ) - self.work_tile_type = ( - SwapAbFc12WorkTileInfo if self.is_swap_ab else NonSwapAbFc12WorkTileInfo - ) + self.launch_cluster_cnt_merge_as_preferred = self.non_clc_mixed_cga_config.launch_cluster_cnt_merge_as_preferred + self.compute_cluster_count = impl_desc.get("compute_cluster_count", self.launch_cluster_cnt_merge_as_preferred) + self.work_tile_type = SwapAbFc12WorkTileInfo if self.is_swap_ab else NonSwapAbFc12WorkTileInfo if self.hint is None: - self.hint = self.launch_cluster_cnt_merge_as_preferred + self.hint = self.compute_cluster_count self.group_hint = self.hint self._validate_configuration() @@ -156,10 +147,14 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.mapping_cluster_shape_mn = self.cluster_shape_mn if self.work_id_mode == "cluster_launch_control": - raise NotImplementedError("cluster_launch_control is not implemented for FC12.") - self._work_id_worker = NonClcMixedCgaSchedulerWorker( - config=self.non_clc_mixed_cga_config, work_id_mode=self.work_id_mode, stream_count=1 - ) + self._work_id_worker = ClcMixedCgaSchedulerWorker(config=self.non_clc_mixed_cga_config) + else: + self._work_id_worker = NonClcMixedCgaSchedulerWorker( + config=self.non_clc_mixed_cga_config, + work_id_mode=self.work_id_mode, + stream_count=1, + work_id_stride=self.compute_cluster_count, + ) def _validate_configuration(self) -> None: if len(self.mma_tiler_mnk) != 3: @@ -181,21 +176,29 @@ def _validate_configuration(self) -> None: raise ValueError("sf_padding_block must be positive.") if self.launch_cluster_cnt_merge_as_preferred <= 0: raise ValueError("launch_cluster_cnt_merge_as_preferred must be positive.") + if ( + isinstance(self.compute_cluster_count, bool) + or not isinstance(self.compute_cluster_count, int) + or not 0 < self.compute_cluster_count <= self.launch_cluster_cnt_merge_as_preferred + ): + raise ValueError( + f"compute_cluster_count must be a Python int in (0, {self.launch_cluster_cnt_merge_as_preferred}], got " + f"{self.compute_cluster_count}." + ) if self.work_id_mode not in ("grid_stride", "atomic_counter", "cluster_launch_control"): raise ValueError( "work_id_mode must be 'grid_stride', 'atomic_counter', or " f"'cluster_launch_control', got {self.work_id_mode!r}." ) cluster_size = self.cluster_shape_mn[0] * self.cluster_shape_mn[1] - if self.work_id_mode == "atomic_counter" and cluster_size > 32: - raise ValueError("The atomic broadcast protocol supports at most 32 CTAs per cluster.") + uses_peer_lane_protocol = self.work_id_mode in ("atomic_counter", "cluster_launch_control") + if uses_peer_lane_protocol and cluster_size > 32: + raise ValueError("The per-lane peer barrier protocol supports at most 32 CTAs per cluster.") fallback_cluster_shape = self.non_clc_mixed_cga_config.fallback_cluster_shape if fallback_cluster_shape is not None: fallback_cluster_size = fallback_cluster_shape[0] * fallback_cluster_shape[1] - if self.work_id_mode == "atomic_counter" and fallback_cluster_size > 32: - raise ValueError( - "The atomic broadcast protocol supports at most 32 CTAs per fallback cluster." - ) + if uses_peer_lane_protocol and fallback_cluster_size > 32: + raise ValueError("The per-lane peer barrier protocol supports at most 32 CTAs per fallback cluster.") for field_name in ("expert_count", "intermediate_gateup_size", "hidden_size"): value = getattr(self, field_name) if isinstance(value, int) and value <= 0: @@ -222,10 +225,48 @@ def initialize_fallback_group(self) -> None: """Register this physical fallback cluster with its fixed logical group.""" self._work_id_worker.initialize_fallback_group() + # CUDA caps gridDim.z at 65535; CLC tickets ride on blockIdx.z. + _clc_grid_z_limit = 65535 + + def clc_ticket_count(self, total_token_rows): + """Return the CLC ticket count, an upper bound on FC12 tasks (surplus tickets decode to DONE).""" + if not isinstance(self.expert_count, int): + raise ValueError("cluster_launch_control needs a static expert_count to size the ticket grid.") + if not isinstance(self.intermediate_gateup_size, int) or not isinstance(self.hidden_size, int): + raise ValueError("cluster_launch_control needs static intermediate/hidden sizes to size the ticket grid.") + if isinstance(total_token_rows, bool) or not isinstance(total_token_rows, int) or total_token_rows < 0: + raise TypeError( + "cluster_launch_control needs a host-known non-negative int total_token_rows; the grid " + "must be validated against the gridDim.z limit before launch." + ) + mapping_cluster_tile_m = self.mapping_cta_tile_shape_mnk[0] * self.mapping_cluster_shape_mn[0] + mapping_cluster_tile_n = self.mapping_cta_tile_shape_mnk[1] * self.mapping_cluster_shape_mn[1] + num_fc1_blocks = (self.intermediate_gateup_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n + num_fc2_blocks = (self.hidden_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n + token_blocks = (total_token_rows + mapping_cluster_tile_m - 1) // mapping_cluster_tile_m + ( + self.expert_count - 1 + ) + # Zero rows still needs one cluster to bootstrap, observe DONE and exit (gridDim.z >= 1). + return max(1, token_blocks * (num_fc1_blocks + num_fc2_blocks)) + def get_grid_shape( - self, *, max_active_clusters: Optional[int] = None, problem_desc=None + self, + *, + max_active_clusters: Optional[int] = None, + problem_desc=None, + total_token_rows=None, ) -> Tuple[int, int, int]: - """Return the persistent launch grid in GEMM-domain orientation.""" + """Return the launch grid in GEMM-domain orientation (full ticket grid under cluster launch control).""" + if self.work_id_mode == "cluster_launch_control": + if total_token_rows is None: + raise ValueError("cluster_launch_control requires total_token_rows to size the ticket grid.") + ticket_count = self.clc_ticket_count(total_token_rows) + if ticket_count > self._clc_grid_z_limit: + raise ValueError( + f"cluster_launch_control ticket count {ticket_count} exceeds gridDim.z limit " + f"{self._clc_grid_z_limit}." + ) + return (self.cluster_shape_mn[0], self.cluster_shape_mn[1], ticket_count) if ( not self.non_clc_mixed_cga_config.is_mixed and max_active_clusters is not None @@ -236,11 +277,7 @@ def get_grid_shape( "launch_cluster_cnt_merge_as_preferred " f"({self.launch_cluster_cnt_merge_as_preferred})." ) - return ( - self.cluster_shape_mn[0], - self.cluster_shape_mn[1], - self.launch_cluster_cnt_merge_as_preferred, - ) + return (self.cluster_shape_mn[0], self.cluster_shape_mn[1], self.launch_cluster_cnt_merge_as_preferred) @cute.jit def assign_device_members( @@ -257,9 +294,7 @@ def assign_device_members( ) -> None: """Initialize all FC12 scheduler state rooted for one CTA lifetime.""" if cutlass.const_expr((expert_token_sizes is None) == (expert_token_prefix_sum is None)): - raise ValueError( - "Exactly one of expert_token_sizes and expert_token_prefix_sum must be provided." - ) + raise ValueError("Exactly one of expert_token_sizes and expert_token_prefix_sum must be provided.") needs_actual_shape = not all( isinstance(dimension, int) for dimension in (self.expert_count, self.intermediate_gateup_size, self.hidden_size) @@ -311,11 +346,30 @@ def gen_next_work(self) -> SchedulerWorkTileBase: cta_id_in_mapping_cluster = _to_fc12_mapping_cta_coord( self._work_id_worker.cta_coord_in_preferred_cluster, self.is_swap_ab ) + if cutlass.const_expr(self.work_id_mode == "cluster_launch_control"): + # CLC tickets on a mixed-cluster launch are not monotonic; the mapper is forward-only, so + # rewind when the ticket falls before the cursor's current segment. + if self._work_id_worker.claimed_is_valid: + if work_id < self._task_mapping_state.cursor_state.current_expert_tile_start: + self._task_mapping_state = rewind_fc12_mapping_cursor(self._task_mapping_state) work_tile, self._task_mapping_state = map_fc12_linear_work_id( work_id, cta_id_in_mapping_cluster, self._task_mapping_state ) + if cutlass.const_expr(self.work_id_mode == "cluster_launch_control"): + # A failed try_cancel means no cluster is left to steal: terminate this CTA's stream. + mapped_tile = work_tile + work_tile = make_fc12_done_tile(self.is_swap_ab) + if self._work_id_worker.claimed_is_valid: + work_tile = mapped_tile return work_tile + @cute.jit + def produce_tail(self) -> None: + """Drain the work transport, then the CLC leader waits for the cluster's final release.""" + super().produce_tail() + if cutlass.const_expr(self.work_id_mode == "cluster_launch_control"): + self._work_id_worker.produce_tail() + def __extract_mlir_values__(self) -> list: values = super().__extract_mlir_values__() values.extend(extract_mlir_values(self._work_id_worker)) @@ -335,9 +389,7 @@ def __new_from_mlir_values__(self, values: list) -> "BlackwellFusedFc12Scheduler def rebuild(state): nonlocal value_index state_value_count = len(extract_mlir_values(state)) - rebuilt_state = new_from_mlir_values( - state, values[value_index : value_index + state_value_count] - ) + rebuilt_state = new_from_mlir_values(state, values[value_index : value_index + state_value_count]) value_index += state_value_count return rebuilt_state @@ -362,6 +414,7 @@ def rebuild(state): "work_id_mode", "is_swap_ab", "launch_cluster_cnt_merge_as_preferred", + "compute_cluster_count", "non_clc_mixed_cga_config", "mapping_cta_tile_shape_mnk", "mapping_cluster_shape_mn", @@ -407,9 +460,7 @@ def __new_from_mlir_values__(self, values: list) -> "_PhaseInterleaveControlStat rebuilt_fields = [] for field in self._fields(): field_value_count = len(extract_mlir_values(field)) - rebuilt_fields.append( - new_from_mlir_values(field, values[value_index : value_index + field_value_count]) - ) + rebuilt_fields.append(new_from_mlir_values(field, values[value_index : value_index + field_value_count])) value_index += field_value_count if value_index != len(values): raise ValueError( @@ -440,7 +491,11 @@ def impl_desc_require(cls) -> dict[str, type]: "sf_padding_block": int, "work_id_mode": str, "is_swap_ab": bool, + "dfc1_m_group": OptionalRequirement(int), + "defer_consumer_until_full_ready": OptionalRequirement(bool), + "fuse_ready_probe_and_linear1_claim": OptionalRequirement(bool), **_mixed_cga_impl_requirements(), + "compute_cluster_count": OptionalRequirement(int), } def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: @@ -458,13 +513,17 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.sf_padding_block = impl_desc["sf_padding_block"] self.work_id_mode: WorkIdAcquisitionMode = impl_desc["work_id_mode"] self.is_swap_ab = impl_desc["is_swap_ab"] - self.non_clc_mixed_cga_config = _make_non_clc_mixed_cga_config(impl_desc) - self.launch_cluster_cnt_merge_as_preferred = ( - self.non_clc_mixed_cga_config.launch_cluster_cnt_merge_as_preferred + self.dfc1_m_group = impl_desc.get("dfc1_m_group", 1) + self.defer_consumer_until_full_ready = impl_desc.get( + "defer_consumer_until_full_ready", False ) - self.work_tile_type = ( - SwapAbFc12WorkTileInfo if self.is_swap_ab else NonSwapAbFc12WorkTileInfo + self.fuse_ready_probe_and_linear1_claim = impl_desc.get( + "fuse_ready_probe_and_linear1_claim", False ) + self.non_clc_mixed_cga_config = _make_non_clc_mixed_cga_config(impl_desc) + self.launch_cluster_cnt_merge_as_preferred = self.non_clc_mixed_cga_config.launch_cluster_cnt_merge_as_preferred + self.compute_cluster_count = impl_desc.get("compute_cluster_count", self.launch_cluster_cnt_merge_as_preferred) + self.work_tile_type = SwapAbFc12WorkTileInfo if self.is_swap_ab else NonSwapAbFc12WorkTileInfo self._validate_configuration() mma_cta_count = 2 if self.use_2cta_instrs else 1 @@ -487,30 +546,39 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: config=self.non_clc_mixed_cga_config, work_id_mode=self.work_id_mode, stream_count=2 ) - mapping_cluster_tile_n = ( - self.mapping_cta_tile_shape_mnk[1] * self.mapping_cluster_shape_mn[1] - ) - self.blocks_fc1 = ( - self.intermediate_gateup_size + mapping_cluster_tile_n - 1 - ) // mapping_cluster_tile_n + mapping_cluster_tile_n = self.mapping_cta_tile_shape_mnk[1] * self.mapping_cluster_shape_mn[1] + self.blocks_fc1 = (self.intermediate_gateup_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n self.blocks_fc2 = (self.hidden_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n - interleave_gcd = math.gcd(self.blocks_fc1, self.blocks_fc2) - self.interleave_fc2_slots = self.blocks_fc2 // interleave_gcd - self.interleave_cycle_length = (self.blocks_fc1 + self.blocks_fc2) // interleave_gcd + self.fc2_claims_per_token_block = self.blocks_fc2 + interleave_gcd = math.gcd( + self.blocks_fc1, self.fc2_claims_per_token_block + ) + self.interleave_fc2_slots = ( + self.fc2_claims_per_token_block // interleave_gcd + ) + self.interleave_cycle_length = ( + self.blocks_fc1 + self.fc2_claims_per_token_block + ) // interleave_gcd self.minimum_global_fc1_claims = minimum_phase_interleave_fc1_claims( blocks_fc1=self.blocks_fc1, blocks_fc2=self.blocks_fc2, - launch_cluster_cnt_merge_as_preferred=self.launch_cluster_cnt_merge_as_preferred, + launch_cluster_cnt_merge_as_preferred=self.compute_cluster_count, ) minimum_hint = minimum_phase_interleave_hint( blocks_fc1=self.blocks_fc1, blocks_fc2=self.blocks_fc2, - launch_cluster_cnt_merge_as_preferred=self.launch_cluster_cnt_merge_as_preferred, + launch_cluster_cnt_merge_as_preferred=self.compute_cluster_count, ) - if self.fc1_prologue_tiles < minimum_hint: + if ( + self.fc1_prologue_tiles < minimum_hint + and not ( + self.fc1_prologue_tiles == 0 + and self.defer_consumer_until_full_ready + ) + ): raise ValueError( f"phase_interleave hint {self.fc1_prologue_tiles} cannot cover a " - f"{self.launch_cluster_cnt_merge_as_preferred}-cluster FC2 claim wave; " + f"{self.compute_cluster_count}-cluster FC2 claim wave; " f"raise hint to at least {minimum_hint}." ) @@ -519,14 +587,10 @@ def _validate_configuration(self) -> None: raise ValueError("mma_tiler_mnk must contain three dimensions.") if len(self.cluster_shape_mn) != 2: raise ValueError("cluster_shape_mn must contain two dimensions.") - if not all( - isinstance(dimension, int) and not isinstance(dimension, bool) - for dimension in self.mma_tiler_mnk - ): + if not all(isinstance(dimension, int) and not isinstance(dimension, bool) for dimension in self.mma_tiler_mnk): raise TypeError("mma_tiler_mnk dimensions must be Python ints.") if not all( - isinstance(dimension, int) and not isinstance(dimension, bool) - for dimension in self.cluster_shape_mn + isinstance(dimension, int) and not isinstance(dimension, bool) for dimension in self.cluster_shape_mn ): raise TypeError("cluster_shape_mn dimensions must be Python ints.") if any(dimension <= 0 for dimension in self.mma_tiler_mnk): @@ -539,21 +603,61 @@ def _validate_configuration(self) -> None: if ( isinstance(self.fc1_prologue_tiles, bool) or not isinstance(self.fc1_prologue_tiles, int) - or self.fc1_prologue_tiles <= 0 + or self.fc1_prologue_tiles < 0 ): raise ValueError( - "fc1_prologue_tiles must be a positive Python int resolved by the kernel frontend." + "fc1_prologue_tiles must be a non-negative Python int " + "resolved by the kernel frontend." ) + if type(self.dfc1_m_group) is not int or self.dfc1_m_group not in (1, 16): + raise ValueError("dfc1_m_group must be 1 or 16") + if self.dfc1_m_group > 1 and ( + self.is_swap_ab + or not self.defer_consumer_until_full_ready + ): + raise ValueError( + "Grouped dFC1 M traversal requires non-swap singleton consumers, " + "ordinary producer order and full-ready admission" + ) + if not isinstance(self.defer_consumer_until_full_ready, bool): + raise TypeError("defer_consumer_until_full_ready must be a bool.") + if not isinstance(self.fuse_ready_probe_and_linear1_claim, bool): + raise TypeError( + "fuse_ready_probe_and_linear1_claim must be a bool." + ) + if ( + self.defer_consumer_until_full_ready + and self.non_clc_mixed_cga_config.is_mixed + ): + raise NotImplementedError( + "defer_consumer_until_full_ready does not yet support fixed " + "preferred/fallback cluster groups." + ) + if self.fuse_ready_probe_and_linear1_claim: + if not self.defer_consumer_until_full_ready: + raise ValueError( + "fuse_ready_probe_and_linear1_claim requires " + "defer_consumer_until_full_ready=True." + ) + if self.fc1_prologue_tiles == 0 and not self.defer_consumer_until_full_ready: + raise ValueError("fc1_prologue_tiles=0 requires deferred full-readiness.") if self.token_padding_block <= 0: raise ValueError("token_padding_block must be positive.") if self.sf_padding_block <= 0: raise ValueError("sf_padding_block must be positive.") if self.launch_cluster_cnt_merge_as_preferred <= 0: raise ValueError("launch_cluster_cnt_merge_as_preferred must be positive.") - if self.work_id_mode != "atomic_counter": + if ( + isinstance(self.compute_cluster_count, bool) + or not isinstance(self.compute_cluster_count, int) + or not 0 < self.compute_cluster_count <= self.launch_cluster_cnt_merge_as_preferred + ): raise ValueError( - "Phase-interleaved FC12 scheduling currently requires work_id_mode='atomic_counter'." + f"compute_cluster_count must be a Python int in (0, {self.launch_cluster_cnt_merge_as_preferred}], got " + f"{self.compute_cluster_count}." ) + if self.work_id_mode != "atomic_counter": + raise ValueError("Phase-interleaved FC12 scheduling currently requires work_id_mode='atomic_counter'.") cluster_size = self.cluster_shape_mn[0] * self.cluster_shape_mn[1] if cluster_size > 32: raise ValueError("The atomic broadcast protocol supports at most 32 CTAs per cluster.") @@ -561,9 +665,7 @@ def _validate_configuration(self) -> None: if fallback_cluster_shape is not None: fallback_cluster_size = fallback_cluster_shape[0] * fallback_cluster_shape[1] if fallback_cluster_size > 32: - raise ValueError( - "The atomic broadcast protocol supports at most 32 CTAs per fallback cluster." - ) + raise ValueError("The atomic broadcast protocol supports at most 32 CTAs per fallback cluster.") maximum_int32 = (1 << 31) - 1 for field_name in ("expert_count", "intermediate_gateup_size", "hidden_size"): value = getattr(self, field_name) @@ -587,7 +689,10 @@ def initialize_fallback_group(self) -> None: self._work_id_worker.initialize_fallback_group() def get_grid_shape( - self, *, max_active_clusters: Optional[int] = None, problem_desc=None + self, + *, + max_active_clusters: Optional[int] = None, + problem_desc=None, ) -> Tuple[int, int, int]: """Return the statically configured persistent launch grid.""" if ( @@ -600,11 +705,7 @@ def get_grid_shape( "launch_cluster_cnt_merge_as_preferred " f"({self.launch_cluster_cnt_merge_as_preferred})." ) - return ( - self.cluster_shape_mn[0], - self.cluster_shape_mn[1], - self.launch_cluster_cnt_merge_as_preferred, - ) + return (self.cluster_shape_mn[0], self.cluster_shape_mn[1], self.launch_cluster_cnt_merge_as_preferred) @cute.jit def assign_device_members( @@ -621,9 +722,7 @@ def assign_device_members( ) -> None: """Initialize phase-local mapping, cadence, and counter state.""" if cutlass.const_expr((expert_token_sizes is None) == (expert_token_prefix_sum is None)): - raise ValueError( - "Exactly one of expert_token_sizes and expert_token_prefix_sum must be provided." - ) + raise ValueError("Exactly one of expert_token_sizes and expert_token_prefix_sum must be provided.") self.create_scheduler_pipelines(smem_workspace, smem_base) self._work_id_worker.assign_device_members( is_fallback_cluster=is_fallback_cluster, @@ -643,14 +742,153 @@ def assign_device_members( is_swap_ab=self.is_swap_ab, expert_token_sizes=expert_token_sizes, expert_token_prefix_sum=expert_token_prefix_sum, + dfc1_m_group=self.dfc1_m_group, ) self._control_state = _PhaseInterleaveControlState( prologue_remaining=Int32(self.fc1_prologue_tiles), - cycle_position=Int32(0), + # Stagger the steady-state cadence by physical cluster. Starting + # every persistent cluster at slot zero turns an intended rolling + # mix into chip-wide FC1/FC2 bursts; the phase offset keeps both + # streams represented in every claim wave after the prologue. + cycle_position=( + Int32(block_idx[2]) % Int32(self.interleave_cycle_length) + ), fc1_exhausted=Boolean(False), fc2_exhausted=Boolean(False), - prologue_claims_synchronized=Boolean(False), + prologue_claims_synchronized=Boolean( + self.defer_consumer_until_full_ready + ), + ) + + @cute.jit + def cluster_uniform_counter_ready( + self, + counter_pointer: cute.Pointer, + ready_threshold: Int32, + ready_is_mask: bool = False, + ) -> Boolean: + """Observe one counter on the cluster leader and broadcast readiness.""" + if cutlass.const_expr(not self.defer_consumer_until_full_ready): + raise ValueError( + "cluster readiness polling requires " + "defer_consumer_until_full_ready=True." + ) + work_id_worker = self._work_id_worker + ready = work_id_worker.cluster_uniform_counter_ready( + counter_pointer, ready_threshold, ready_is_mask ) + self._work_id_worker = work_id_worker + return ready + + @cute.jit + def gen_next_linear1_work_unless_counter_ready( + self, + counter_pointer: cute.Pointer, + ready_threshold: Int32, + ready_is_mask: bool = False, + ) -> Tuple[Boolean, SchedulerWorkTileBase]: + """Probe full readiness and otherwise claim one Linear1 tile. + + When Linear1 still has unclaimed work, the readiness load and atomic + claim share one cluster broadcast transaction. Once the stream is + exhausted this falls back to a probe-only transaction, avoiding + repeated tail over-claims while the held consumer drains. + """ + if cutlass.const_expr(not self.fuse_ready_probe_and_linear1_claim): + raise ValueError( + "fused readiness/Linear1 claims require " + "fuse_ready_probe_and_linear1_claim=True." + ) + + work_tile = make_fc12_done_tile(self.is_swap_ab) + work_id_worker = self._work_id_worker + task_mapping_state = self._task_mapping_state + control_state = self._control_state + fc1_exhausted = control_state.fc1_exhausted + ready = Boolean(False) + + if fc1_exhausted: + ready = work_id_worker.cluster_uniform_counter_ready( + counter_pointer, ready_threshold, ready_is_mask + ) + else: + ready, linear_work_id = ( + work_id_worker.cluster_uniform_counter_ready_or_claim( + counter_pointer, + ready_threshold, + claim_stream_index=Int32(0), + ready_is_mask=ready_is_mask, + ) + ) + if not ready: + cta_id_in_mapping_cluster = _to_fc12_mapping_cta_coord( + work_id_worker.cta_coord_in_preferred_cluster, + self.is_swap_ab, + ) + work_tile, stream_has_work, task_mapping_state = ( + map_phase_interleaved_fc12_work_id( + linear_work_id, + Int32(BlockPhase.Linear1), + cta_id_in_mapping_cluster, + task_mapping_state, + ) + ) + if not stream_has_work: + fc1_exhausted = Boolean(True) + + control_state.fc1_exhausted = fc1_exhausted + self._work_id_worker = work_id_worker + self._task_mapping_state = task_mapping_state + self._control_state = control_state + return ready, work_tile + + + @cute.jit + def gen_next_linear1_work(self) -> SchedulerWorkTileBase: + """Force one dependency-free Linear1 claim while a Linear2 tile is held. + + This corrective claim is outside the normal rolling cadence: the + deferred Linear2 claim already consumed its cadence slot. The caller + must publish a returned Linear1 tile before rechecking or publishing + the held consumer. Consequently no producer descriptor can become + trapped behind an unready consumer. + """ + if cutlass.const_expr(not self.defer_consumer_until_full_ready): + raise ValueError( + "forced Linear1 claims require " + "defer_consumer_until_full_ready=True." + ) + + work_tile = make_fc12_done_tile(self.is_swap_ab) + work_id_worker = self._work_id_worker + task_mapping_state = self._task_mapping_state + control_state = self._control_state + fc1_exhausted = control_state.fc1_exhausted + + if not fc1_exhausted: + linear_work_id = work_id_worker.claim_next_work( + Int32(0) + ) + cta_id_in_mapping_cluster = _to_fc12_mapping_cta_coord( + work_id_worker.cta_coord_in_preferred_cluster, + self.is_swap_ab, + ) + work_tile, stream_has_work, task_mapping_state = ( + map_phase_interleaved_fc12_work_id( + linear_work_id, + Int32(BlockPhase.Linear1), + cta_id_in_mapping_cluster, + task_mapping_state, + ) + ) + if not stream_has_work: + fc1_exhausted = Boolean(True) + + control_state.fc1_exhausted = fc1_exhausted + self._work_id_worker = work_id_worker + self._task_mapping_state = task_mapping_state + self._control_state = control_state + return work_tile @cute.jit def _wait_for_fc1_prologue_claims( @@ -664,10 +902,8 @@ def _wait_for_fc1_prologue_claims( fc1_counter_pointer = work_id_worker.get_atomic_counter_pointer(Int32(0)) minimum_claim_count = Int32(self.minimum_global_fc1_claims) if not spin_peek(fc1_counter_pointer, lambda value: value >= minimum_claim_count): - claim_target, stream_ends_before_target = ( - resolve_phase_interleaved_fc1_claim_target( - minimum_claim_count, task_mapping_state - ) + claim_target, stream_ends_before_target = resolve_phase_interleaved_fc1_claim_target( + minimum_claim_count, task_mapping_state ) spin_wait(fc1_counter_pointer, lambda value: value >= claim_target) return fc1_exhausted @@ -683,6 +919,7 @@ def gen_next_work(self) -> SchedulerWorkTileBase: cycle_position = control_state.cycle_position fc1_exhausted = control_state.fc1_exhausted fc2_exhausted = control_state.fc2_exhausted + stream_has_work = Boolean(False) prologue_claims_synchronized = control_state.prologue_claims_synchronized resolved = Boolean(False) @@ -690,7 +927,8 @@ def gen_next_work(self) -> SchedulerWorkTileBase: if fc1_exhausted and fc2_exhausted: work_tile = make_fc12_done_tile(self.is_swap_ab) resolved = Boolean(True) - else: + + if not resolved: want_fc1 = Boolean(True) if prologue_remaining <= Int32(0): is_fc2_slot = (cycle_position * Int32(self.interleave_fc2_slots)) % Int32( @@ -707,14 +945,19 @@ def gen_next_work(self) -> SchedulerWorkTileBase: atomic_counter_index = Int32(0) linear_work_id = work_id_worker.claim_next_work(atomic_counter_index) want_fc1 = work_id_worker.claimed_stream_index == Int32(0) - phase = Int32(BlockPhase.Linear2) - if want_fc1: - phase = Int32(BlockPhase.Linear1) cta_id_in_mapping_cluster = _to_fc12_mapping_cta_coord( work_id_worker.cta_coord_in_preferred_cluster, self.is_swap_ab ) - work_tile, stream_has_work, task_mapping_state = map_phase_interleaved_fc12_work_id( - linear_work_id, phase, cta_id_in_mapping_cluster, task_mapping_state + phase = Int32(BlockPhase.Linear2) + if want_fc1: + phase = Int32(BlockPhase.Linear1) + work_tile, stream_has_work, task_mapping_state = ( + map_phase_interleaved_fc12_work_id( + linear_work_id, + phase, + cta_id_in_mapping_cluster, + task_mapping_state, + ) ) if stream_has_work: if (not want_fc1) and (not prologue_claims_synchronized): @@ -725,9 +968,7 @@ def gen_next_work(self) -> SchedulerWorkTileBase: if prologue_remaining > Int32(0): prologue_remaining = prologue_remaining - Int32(1) else: - cycle_position = (cycle_position + Int32(1)) % Int32( - self.interleave_cycle_length - ) + cycle_position = (cycle_position + Int32(1)) % Int32(self.interleave_cycle_length) resolved = Boolean(True) else: if want_fc1: @@ -764,9 +1005,7 @@ def __new_from_mlir_values__(self, values: list) -> "PhaseInterleavedFc12Schedul def rebuild(state): nonlocal value_index state_value_count = len(extract_mlir_values(state)) - rebuilt_state = new_from_mlir_values( - state, values[value_index : value_index + state_value_count] - ) + rebuilt_state = new_from_mlir_values(state, values[value_index : value_index + state_value_count]) value_index += state_value_count return rebuilt_state @@ -792,11 +1031,16 @@ def rebuild(state): "work_id_mode", "is_swap_ab", "launch_cluster_cnt_merge_as_preferred", + "compute_cluster_count", "non_clc_mixed_cga_config", "mapping_cta_tile_shape_mnk", "mapping_cluster_shape_mn", "blocks_fc1", "blocks_fc2", + "dfc1_m_group", + "defer_consumer_until_full_ready", + "fuse_ready_probe_and_linear1_claim", + "fc2_claims_per_token_block", "interleave_fc2_slots", "interleave_cycle_length", "minimum_global_fc1_claims", @@ -805,8 +1049,4 @@ def rebuild(state): return result -__all__ = [ - "BlackwellFusedFc12Scheduler", - "PhaseInterleavedFc12Scheduler", - "minimum_phase_interleave_hint", -] +__all__ = ["BlackwellFusedFc12Scheduler", "PhaseInterleavedFc12Scheduler", "minimum_phase_interleave_hint"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/non_clc_mixed_cga.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/non_clc_mixed_cga.py index 90f340877572..ba879cc6ac4d 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/non_clc_mixed_cga.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/non_clc_mixed_cga.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Reusable preferred/fallback cluster scheduling without hardware CLC.""" import dataclasses @@ -19,6 +16,8 @@ AtomicCounterWorkIdState, FixedGroupMixedCgaAtomicCounterWorkIdState, GridStrideWorkIdState, + broadcast_atomic_counter_ready, + broadcast_atomic_counter_ready_or_claim, claim_work_id, initialize_fixed_group_mixed_cga_work_id_state, ) @@ -49,19 +48,14 @@ def __post_init__(self) -> None: and not isinstance(self.fallback_cluster_count, bool) and self.fallback_cluster_count == 0 ) - if ( - self.fallback_cluster_shape == self.preferred_cluster_shape - or has_no_fallback_clusters - ): + if self.fallback_cluster_shape == self.preferred_cluster_shape or has_no_fallback_clusters: object.__setattr__(self, "fallback_cluster_shape", None) object.__setattr__(self, "preferred_cluster_count", None) object.__setattr__(self, "fallback_cluster_count", None) if self.fallback_cluster_shape is None: if not self._is_positive_int(self.launch_cluster_count): - raise ValueError( - "launch_cluster_count must be positive when fallback_cluster_shape is absent." - ) + raise ValueError("launch_cluster_count must be positive when fallback_cluster_shape is absent.") mn_split_factors = (1, 1) split_factor = 1 launch_cluster_cnt_merge_as_preferred = self.launch_cluster_count @@ -70,9 +64,7 @@ def __post_init__(self) -> None: if len(self.fallback_cluster_shape) != len(self.preferred_cluster_shape): raise ValueError("preferred and fallback cluster shapes must have equal rank.") if not self._is_positive_int(self.preferred_cluster_count): - raise ValueError( - "preferred_cluster_count must be positive when fallback_cluster_shape is present." - ) + raise ValueError("preferred_cluster_count must be positive when fallback_cluster_shape is present.") if ( isinstance(self.fallback_cluster_count, bool) or not isinstance(self.fallback_cluster_count, int) @@ -92,16 +84,12 @@ def __post_init__(self) -> None: self.preferred_cluster_shape, self.fallback_cluster_shape ) ): - raise ValueError( - "Every preferred cluster dimension must be divisible by its fallback dimension." - ) + raise ValueError("Every preferred cluster dimension must be divisible by its fallback dimension.") split_factor = self._shape_size(mn_split_factors) is_mixed = self.fallback_cluster_shape != self.preferred_cluster_shape if is_mixed: if split_factor > 16 or split_factor & (split_factor - 1): - raise ValueError( - "The preferred/fallback cluster split factor must be a power of two at most 16." - ) + raise ValueError("The preferred/fallback cluster split factor must be a power of two at most 16.") if self.fallback_cluster_count % split_factor != 0: raise ValueError( "fallback_cluster_count must be divisible by the preferred/fallback cluster split factor." @@ -110,22 +98,14 @@ def __post_init__(self) -> None: self.preferred_cluster_count + self.fallback_cluster_count // split_factor ) else: - launch_cluster_cnt_merge_as_preferred = ( - self.preferred_cluster_count + self.fallback_cluster_count - ) + launch_cluster_cnt_merge_as_preferred = self.preferred_cluster_count + self.fallback_cluster_count if launch_cluster_cnt_merge_as_preferred <= 0: - raise ValueError( - "The resolved launch must contain at least one cluster merged as preferred." - ) + raise ValueError("The resolved launch must contain at least one cluster merged as preferred.") object.__setattr__(self, "mn_split_factors", mn_split_factors) object.__setattr__(self, "split_factor", split_factor) - object.__setattr__( - self, "launch_cluster_cnt_merge_as_preferred", launch_cluster_cnt_merge_as_preferred - ) - object.__setattr__( - self, "total_cta_cnt", launch_cluster_cnt_merge_as_preferred * preferred_cluster_size - ) + object.__setattr__(self, "launch_cluster_cnt_merge_as_preferred", launch_cluster_cnt_merge_as_preferred) + object.__setattr__(self, "total_cta_cnt", launch_cluster_cnt_merge_as_preferred * preferred_cluster_size) object.__setattr__(self, "is_mixed", is_mixed) @staticmethod @@ -162,16 +142,22 @@ def __init__( config: NonClcMixedCgaConfig, work_id_mode: WorkIdAcquisitionMode, stream_count: int, + work_id_stride: Optional[int] = None, ) -> None: if work_id_mode not in ("grid_stride", "atomic_counter"): - raise ValueError( - "Non-CLC mixed-CGA scheduling supports grid_stride or atomic_counter work IDs." - ) + raise ValueError("Non-CLC mixed-CGA scheduling supports grid_stride or atomic_counter work IDs.") if isinstance(stream_count, bool) or not isinstance(stream_count, int) or stream_count <= 0: raise ValueError("stream_count must be a positive Python int.") + if work_id_stride is None: + work_id_stride = config.launch_cluster_cnt_merge_as_preferred + if isinstance(work_id_stride, bool) or not isinstance(work_id_stride, int) or work_id_stride <= 0: + raise ValueError("work_id_stride must be a positive Python int.") self.config = config self.work_id_mode = work_id_mode self.stream_count = stream_count + # The work stream can have fewer participants than the physical launch + # when some clusters only communicate. Atomic claims do not use this. + self.work_id_stride = work_id_stride def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: """Register the cluster-local claim broadcast channel.""" @@ -240,9 +226,9 @@ def assign_device_members( for dimension_idx in cutlass.range_constexpr(len(self.config.mn_split_factors)): preferred_dimension = self.config.preferred_cluster_shape[dimension_idx] fallback_dimension = self.config.fallback_cluster_shape[dimension_idx] - inner_coordinate = ( - Int32(block_idx[dimension_idx]) % Int32(preferred_dimension) - ) // Int32(fallback_dimension) + inner_coordinate = (Int32(block_idx[dimension_idx]) % Int32(preferred_dimension)) // Int32( + fallback_dimension + ) flattened_index = flattened_index + inner_coordinate * Int32(dimension_stride) dimension_stride *= self.config.mn_split_factors[dimension_idx] cta_coord_in_preferred_cluster = self._preferred_cluster_cta_coord( @@ -255,20 +241,15 @@ def assign_device_members( cluster_pipeline = pipeline.PipelineAsync.create( num_stages=1, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, 32 * active_cluster_size - ), - barrier_storage=smem_workspace.ptr( - self.cluster_pipeline_mbarriers_region, smem_base - ), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * active_cluster_size), + barrier_storage=smem_workspace.ptr(self.cluster_pipeline_mbarriers_region, smem_base), defer_sync=True, ) atomic_counter_state = AtomicCounterWorkIdState( counter_pointer=device_workspace.ptr(self.work_id_counter_region), counter_count=self.stream_count, broadcast_pointer=smem_workspace.ptr(self.cluster_broadcast_region, smem_base), - is_leader_cta=(cta_coord_in_active_cluster[0] + cta_coord_in_active_cluster[1]) - == Int32(0), + is_leader_cta=(cta_coord_in_active_cluster[0] + cta_coord_in_active_cluster[1]) == Int32(0), cluster_pipeline=cluster_pipeline, producer_state=pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1), consumer_state=pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1), @@ -278,9 +259,7 @@ def assign_device_members( self._cta_coord_in_active_cluster = cta_coord_in_active_cluster self._work_id_state = FixedGroupMixedCgaAtomicCounterWorkIdState( atomic_counter_state=atomic_counter_state, - registration_counter_pointer=device_workspace.ptr( - self.fallback_registration_counter_region - ), + registration_counter_pointer=device_workspace.ptr(self.fallback_registration_counter_region), group_token_pointer=device_workspace.ptr(self.fallback_group_token_region), split_factor=self.config.split_factor, fallback_cluster_count=self.config.fallback_cluster_count, @@ -296,7 +275,7 @@ def assign_device_members( else: self._work_id_state = GridStrideWorkIdState( next_work_id=Int32(block_idx[2]), - work_id_stride=Int32(self.config.launch_cluster_cnt_merge_as_preferred), + work_id_stride=Int32(self.work_id_stride), ) self.claimed_stream_index = Int32(0) @@ -304,9 +283,7 @@ def assign_device_members( def initialize_fallback_group(self) -> None: """Register one physical fallback cluster with its fixed logical group.""" if cutlass.const_expr(self.config.is_mixed and self.work_id_mode == "atomic_counter"): - self._work_id_state = initialize_fixed_group_mixed_cga_work_id_state( - self._work_id_state - ) + self._work_id_state = initialize_fixed_group_mixed_cga_work_id_state(self._work_id_state) if self._work_id_state.is_fallback_cluster: self.cta_coord_in_preferred_cluster = self._preferred_cluster_cta_coord( self._cta_coord_in_active_cluster, self._work_id_state.in_group_idx @@ -321,10 +298,8 @@ def _preferred_cluster_cta_coord( self.config.mn_split_factors[1] ) return ( - cta_coord_in_active_cluster[0] - + inner_cluster_m * Int32(self.config.fallback_cluster_shape[0]), - cta_coord_in_active_cluster[1] - + inner_cluster_n * Int32(self.config.fallback_cluster_shape[1]), + cta_coord_in_active_cluster[0] + inner_cluster_m * Int32(self.config.fallback_cluster_shape[0]), + cta_coord_in_active_cluster[1] + inner_cluster_n * Int32(self.config.fallback_cluster_shape[1]), Int32(0), ) @@ -334,18 +309,14 @@ def get_atomic_counter_pointer(self, stream_index: Int32) -> cute.Pointer: if cutlass.const_expr(self.work_id_mode != "atomic_counter"): raise ValueError("Work-ID counters are only available in atomic_counter mode.") work_id_state = self._work_id_state - if cutlass.const_expr( - isinstance(work_id_state, FixedGroupMixedCgaAtomicCounterWorkIdState) - ): + if cutlass.const_expr(isinstance(work_id_state, FixedGroupMixedCgaAtomicCounterWorkIdState)): work_id_state = work_id_state.atomic_counter_state return work_id_state.counter_pointer + stream_index @cute.jit def claim_next_work(self, stream_index=0) -> Int32: """Claim the next canonical work ID and update the preferred CTA coordinate.""" - claimed_work_id, self._work_id_state = claim_work_id( - self._work_id_state, atomic_counter_index=stream_index - ) + claimed_work_id, self._work_id_state = claim_work_id(self._work_id_state, atomic_counter_index=stream_index) canonical_work_id = claimed_work_id if cutlass.const_expr(self.config.is_mixed and self.work_id_mode == "atomic_counter"): cta_coord_in_preferred_cluster = self._cta_coord_in_active_cluster @@ -359,6 +330,69 @@ def claim_next_work(self, stream_index=0) -> Int32: self.claimed_stream_index = Int32(stream_index) return canonical_work_id + @cute.jit + def cluster_uniform_counter_ready( + self, + counter_pointer: cute.Pointer, + ready_threshold: Int32, + ready_is_mask: bool = False, + ) -> Boolean: + """Return one leader-loaded readiness predicate to every active CTA. + + V1 deliberately excludes fixed fallback groups. Their several + physical clusters must rendezvous for every logical work-ID claim, so + independently polling readiness could make group members take + different claim paths. Homogeneous preferred clusters have one shared + claim/broadcast channel and therefore one uniform decision. + """ + if cutlass.const_expr(self.work_id_mode != "atomic_counter"): + raise ValueError( + "cluster_uniform_counter_ready requires atomic-counter work IDs." + ) + if cutlass.const_expr(self.config.is_mixed): + raise NotImplementedError( + "cluster-uniform readiness polling does not yet support fixed " + "preferred/fallback groups." + ) + ready, self._work_id_state = broadcast_atomic_counter_ready( + self._work_id_state, + counter_pointer, + ready_threshold, + ready_is_mask, + ) + return ready + + @cute.jit + def cluster_uniform_counter_ready_or_claim( + self, + ready_counter_pointer: cute.Pointer, + ready_threshold: Int32, + claim_stream_index=0, + ready_is_mask: bool = False, + ) -> Tuple[Boolean, Int32]: + """Probe readiness or claim one ID, uniformly across one cluster.""" + if cutlass.const_expr(self.work_id_mode != "atomic_counter"): + raise ValueError( + "cluster_uniform_counter_ready_or_claim requires " + "atomic-counter work IDs." + ) + if cutlass.const_expr(self.config.is_mixed): + raise NotImplementedError( + "fused cluster readiness/claim does not support fixed " + "preferred/fallback groups." + ) + ready, claimed_work_id, self._work_id_state = ( + broadcast_atomic_counter_ready_or_claim( + self._work_id_state, + ready_counter_pointer, + ready_threshold, + atomic_counter_index=claim_stream_index, + ready_is_mask=ready_is_mask, + ) + ) + self.claimed_stream_index = Int32(claim_stream_index) + return ready, claimed_work_id + def __extract_mlir_values__(self) -> List[ir.Value]: values: List[ir.Value] = [] values.extend(extract_mlir_values(self._work_id_state)) @@ -374,14 +408,15 @@ def __new_from_mlir_values__(self, values: List[ir.Value]) -> "NonClcMixedCgaSch def rebuild(field): nonlocal value_index field_value_count = len(extract_mlir_values(field)) - result = new_from_mlir_values( - field, values[value_index : value_index + field_value_count] - ) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) value_index += field_value_count return result result = type(self)( - config=self.config, work_id_mode=self.work_id_mode, stream_count=self.stream_count + config=self.config, + work_id_mode=self.work_id_mode, + stream_count=self.stream_count, + work_id_stride=self.work_id_stride, ) result._work_id_state = rebuild(self._work_id_state) result.cta_coord_in_preferred_cluster = rebuild(self.cta_coord_in_preferred_cluster) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/work_id_claim.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/work_id_claim.py index f9e7b799aed5..846fbe346061 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/work_id_claim.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/work_id_claim.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Blackwell persistent work-ID claim backends.""" import dataclasses @@ -10,13 +7,10 @@ import cutlass.cute as cute import cutlass.pipeline as pipeline from cutlass._mlir import ir +from cutlass._mlir.dialects import nvvm as _nvvm from cutlass.cutlass_dsl import Boolean, Int32, Int64, extract_mlir_values, new_from_mlir_values -from ...helpers.ptx_helpers import ( - mbarrier_arrive_expect_tx_on_peer, - nanosleep, - store_i32_to_peer_cluster_smem_async, -) +from ...helpers.ptx_helpers import mbarrier_arrive_expect_tx_on_peer, nanosleep, store_i32_to_peer_cluster_smem_async class GridStrideWorkIdState: @@ -42,9 +36,7 @@ def __new_from_mlir_values__(self, values: List[ir.Value]) -> "GridStrideWorkIdS ) return type(self)( next_work_id=new_from_mlir_values(self.next_work_id, values[:next_work_id_value_count]), - work_id_stride=new_from_mlir_values( - self.work_id_stride, values[next_work_id_value_count:] - ), + work_id_stride=new_from_mlir_values(self.work_id_stride, values[next_work_id_value_count:]), ) @@ -62,11 +54,7 @@ def __init__( consumer_state, cluster_size: int | Int32, ) -> None: - if ( - isinstance(counter_count, bool) - or not isinstance(counter_count, int) - or counter_count <= 0 - ): + if isinstance(counter_count, bool) or not isinstance(counter_count, int) or counter_count <= 0: raise ValueError("counter_count must be a positive Python int.") self.counter_pointer = counter_pointer self.counter_count = counter_count @@ -97,9 +85,7 @@ def __new_from_mlir_values__(self, values: List[ir.Value]) -> "AtomicCounterWork def rebuild(field): nonlocal value_index field_value_count = len(extract_mlir_values(field)) - result = new_from_mlir_values( - field, values[value_index : value_index + field_value_count] - ) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) value_index += field_value_count return result @@ -111,9 +97,7 @@ def rebuild(field): cluster_pipeline=self.cluster_pipeline, producer_state=rebuild(self.producer_state), consumer_state=rebuild(self.consumer_state), - cluster_size=rebuild(self.cluster_size) - if isinstance(self.cluster_size, Int32) - else self.cluster_size, + cluster_size=rebuild(self.cluster_size) if isinstance(self.cluster_size, Int32) else self.cluster_size, ) if value_index != len(values): raise ValueError( @@ -180,17 +164,13 @@ def __extract_mlir_values__(self) -> List[ir.Value]: values.extend(extract_mlir_values(field)) return values - def __new_from_mlir_values__( - self, values: List[ir.Value] - ) -> "FixedGroupMixedCgaAtomicCounterWorkIdState": + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "FixedGroupMixedCgaAtomicCounterWorkIdState": value_index = 0 def rebuild(field): nonlocal value_index field_value_count = len(extract_mlir_values(field)) - result = new_from_mlir_values( - field, values[value_index : value_index + field_value_count] - ) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) value_index += field_value_count return result @@ -234,9 +214,7 @@ def __new_from_mlir_values__(self, values: List[ir.Value]) -> "GridWorkId": if len(values) != 4: raise ValueError(f"GridWorkId expects four MLIR values, got {len(values)}.") fields = (self.grid_m, self.grid_n, self.grid_l, self.is_valid) - return type(self)( - *(new_from_mlir_values(field, [value]) for field, value in zip(fields, values)) - ) + return type(self)(*(new_from_mlir_values(field, [value]) for field, value in zip(fields, values))) class ClusterLaunchControlWorkIdState: @@ -291,9 +269,7 @@ def __new_from_mlir_values__(self, values: List[ir.Value]) -> "ClusterLaunchCont def rebuild(field): nonlocal value_index field_value_count = len(extract_mlir_values(field)) - rebuilt_field = new_from_mlir_values( - field, values[value_index : value_index + field_value_count] - ) + rebuilt_field = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) value_index += field_value_count return rebuilt_field @@ -317,10 +293,93 @@ def rebuild(field): return result +class ClcMixedCgaWorkIdState: + """Cluster-wide state for hardware CLC claims on a runtime-sized (mixed-CGA) cluster. + + ``PipelineClcFetchAsync`` needs a static cluster layout, so this drives a plain ``PipelineAsync`` + with the atomic broadcast's peer barrier protocol instead. + """ + + def __init__( + self, + response_pending: Boolean, + grid_m: Int32, + grid_n: Int32, + grid_l: Int32, + response_is_valid: Boolean, + cta_coord_in_cluster: cute.Coord, + cluster_pipeline: pipeline.PipelineAsync, + producer_state, + consumer_state, + is_leader_cta: Boolean, + response_pointer: cute.Pointer, + cluster_size: int | Int32, + ) -> None: + self.response_pending = response_pending + self.grid_m = grid_m + self.grid_n = grid_n + self.grid_l = grid_l + self.response_is_valid = response_is_valid + self.cta_coord_in_cluster = cta_coord_in_cluster + self.cluster_pipeline = cluster_pipeline + self.producer_state = producer_state + self.consumer_state = consumer_state + self.is_leader_cta = is_leader_cta + self.response_pointer = response_pointer + self.cluster_size = cluster_size + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + for field in ( + self.response_pending, + self.grid_m, + self.grid_n, + self.grid_l, + self.response_is_valid, + self.cta_coord_in_cluster, + self.producer_state, + self.consumer_state, + self.is_leader_cta, + self.response_pointer, + ): + values.extend(extract_mlir_values(field)) + if isinstance(self.cluster_size, Int32): + values.extend(extract_mlir_values(self.cluster_size)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "ClcMixedCgaWorkIdState": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + rebuilt_field = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) + value_index += field_value_count + return rebuilt_field + + result = type(self)( + response_pending=rebuild(self.response_pending), + grid_m=rebuild(self.grid_m), + grid_n=rebuild(self.grid_n), + grid_l=rebuild(self.grid_l), + response_is_valid=rebuild(self.response_is_valid), + cta_coord_in_cluster=rebuild(self.cta_coord_in_cluster), + cluster_pipeline=self.cluster_pipeline, + producer_state=rebuild(self.producer_state), + consumer_state=rebuild(self.consumer_state), + is_leader_cta=rebuild(self.is_leader_cta), + response_pointer=rebuild(self.response_pointer), + cluster_size=rebuild(self.cluster_size) if isinstance(self.cluster_size, Int32) else self.cluster_size, + ) + if value_index != len(values): + raise ValueError( + f"ClcMixedCgaWorkIdState MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return result + + @cute.jit -def _claim_grid_stride_work_id( - work_id_state: GridStrideWorkIdState, -) -> Tuple[Int32, GridStrideWorkIdState]: +def _claim_grid_stride_work_id(work_id_state: GridStrideWorkIdState) -> Tuple[Int32, GridStrideWorkIdState]: """Claim the next ID from one monotonic grid-stride stream.""" linear_work_id = work_id_state.next_work_id work_id_state.next_work_id = linear_work_id + work_id_state.work_id_stride @@ -345,17 +404,13 @@ def _claim_atomic_counter_work_id( if work_id_state.is_leader_cta: cluster_pipeline.producer_acquire(work_id_state.producer_state) - full_barrier_pointer = cluster_pipeline.sync_object_full.get_barrier( - work_id_state.producer_state.index - ) + full_barrier_pointer = cluster_pipeline.sync_object_full.get_barrier(work_id_state.producer_state.index) thread_idx, _, _ = cute.arch.thread_idx() lane_idx = thread_idx % Int32(32) atomic_work_id = Int32(0) if lane_idx == Int32(0): atomic_work_id = cute.arch.atomic_add(selected_counter_pointer, Int32(1)) - atomic_work_id = cute.arch.shuffle_sync( - atomic_work_id, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31 - ) + atomic_work_id = cute.arch.shuffle_sync(atomic_work_id, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31) if lane_idx < Int32(work_id_state.cluster_size): store_i32_to_peer_cluster_smem_async( work_id_state.broadcast_pointer, atomic_work_id, full_barrier_pointer, lane_idx @@ -371,6 +426,165 @@ def _claim_atomic_counter_work_id( return linear_work_id, work_id_state +@cute.jit +def broadcast_atomic_counter_ready( + work_id_state: AtomicCounterWorkIdState, + counter_pointer: cute.Pointer, + ready_threshold: Int32, + ready_is_mask: bool = False, +) -> Tuple[Boolean, AtomicCounterWorkIdState]: + """Load one global readiness counter on the leader and broadcast the predicate. + + The scheduler makes publication decisions at preferred-cluster granularity, + so a CTA-local peek is not sufficient: the two scheduler warps must either + both defer or both publish their CTA halves of a work tile. Reusing the + atomic work-ID broadcast channel gives the readiness decision the same + cluster-uniform ordering as a work-ID claim. + """ + if cutlass.const_expr(not isinstance(ready_is_mask, bool)): + raise TypeError("ready_is_mask must be a Python bool.") + + broadcast_tensor = cute.make_tensor( + work_id_state.broadcast_pointer, cute.make_layout((1,)) + ) + cluster_pipeline = work_id_state.cluster_pipeline + + if work_id_state.is_leader_cta: + cluster_pipeline.producer_acquire(work_id_state.producer_state) + full_barrier_pointer = cluster_pipeline.sync_object_full.get_barrier( + work_id_state.producer_state.index + ) + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + ready_payload = Int32(0) + if lane_idx == Int32(0): + observed = cute.arch.load( + counter_pointer, Int32, sem="acquire", scope="gpu" + ) + if cutlass.const_expr(ready_is_mask): + if (observed & ready_threshold) == ready_threshold: + ready_payload = Int32(1) + else: + if observed >= ready_threshold: + ready_payload = Int32(1) + ready_payload = cute.arch.shuffle_sync( + ready_payload, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31 + ) + if lane_idx < Int32(work_id_state.cluster_size): + store_i32_to_peer_cluster_smem_async( + work_id_state.broadcast_pointer, + ready_payload, + full_barrier_pointer, + lane_idx, + ) + mbarrier_arrive_expect_tx_on_peer( + full_barrier_pointer, Int32(4), lane_idx + ) + work_id_state.producer_state.advance() + + cluster_pipeline.consumer_wait(work_id_state.consumer_state) + ready_payload = broadcast_tensor[0] + cute.arch.fence_acq_rel_cta() + cluster_pipeline.sync_object_empty.arrive( + work_id_state.consumer_state.index, Int32(0) + ) + work_id_state.consumer_state.advance() + return ready_payload != Int32(0), work_id_state + + +@cute.jit +def broadcast_atomic_counter_ready_or_claim( + work_id_state: AtomicCounterWorkIdState, + ready_counter_pointer: cute.Pointer, + ready_threshold: Int32, + atomic_counter_index=0, + ready_is_mask: bool = False, +) -> Tuple[Boolean, Int32, AtomicCounterWorkIdState]: + """Probe readiness or claim one work ID in one cluster transaction. + + The leader first performs the same acquire readiness load as + :func:`broadcast_atomic_counter_ready`. When the counter is not ready it + immediately claims one ID from ``atomic_counter_index`` before broadcasting + the result. A negative payload denotes readiness; valid work IDs are + non-negative, so the one-word cluster broadcast ABI remains unchanged. + + This primitive is intentionally limited to a single claimed ID. It is the + bundle-1 fast path used while a full-readiness-gated consumer is held + outside the execution FIFO. + """ + if cutlass.const_expr(not isinstance(ready_is_mask, bool)): + raise TypeError("ready_is_mask must be a Python bool.") + invalid_static_index = isinstance(atomic_counter_index, int) and ( + atomic_counter_index < 0 + or atomic_counter_index >= work_id_state.counter_count + ) + if cutlass.const_expr(invalid_static_index): + raise ValueError( + "atomic_counter_index must be in [0, " + f"{work_id_state.counter_count}), got {atomic_counter_index}." + ) + + broadcast_tensor = cute.make_tensor( + work_id_state.broadcast_pointer, cute.make_layout((1,)) + ) + cluster_pipeline = work_id_state.cluster_pipeline + selected_counter_pointer = ( + work_id_state.counter_pointer + Int32(atomic_counter_index) + ) + + if work_id_state.is_leader_cta: + cluster_pipeline.producer_acquire(work_id_state.producer_state) + full_barrier_pointer = cluster_pipeline.sync_object_full.get_barrier( + work_id_state.producer_state.index + ) + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + # -1 is the ready sentinel. A failed probe replaces it with the + # non-negative ID returned by the Linear1 atomic counter. + payload = Int32(-1) + if lane_idx == Int32(0): + observed = cute.arch.load( + ready_counter_pointer, Int32, sem="acquire", scope="gpu" + ) + ready = Boolean(False) + if cutlass.const_expr(ready_is_mask): + ready = (observed & ready_threshold) == ready_threshold + else: + ready = observed >= ready_threshold + if not ready: + payload = cute.arch.atomic_add( + selected_counter_pointer, Int32(1) + ) + payload = Int32( + cute.arch.shuffle_sync( + payload, + offset=0, + mask=0xFFFFFFFF, + mask_and_clamp=31, + ) + ) + if lane_idx < Int32(work_id_state.cluster_size): + store_i32_to_peer_cluster_smem_async( + work_id_state.broadcast_pointer, + payload, + full_barrier_pointer, + lane_idx, + ) + mbarrier_arrive_expect_tx_on_peer( + full_barrier_pointer, Int32(4), lane_idx + ) + work_id_state.producer_state.advance() + + cluster_pipeline.consumer_wait(work_id_state.consumer_state) + payload = broadcast_tensor[0] + cute.arch.fence_acq_rel_cta() + cluster_pipeline.sync_object_empty.arrive( + work_id_state.consumer_state.index, Int32(0) + ) + work_id_state.consumer_state.advance() + return payload == Int32(-1), payload, work_id_state + + @cute.jit def initialize_fixed_group_mixed_cga_work_id_state( work_id_state: FixedGroupMixedCgaAtomicCounterWorkIdState, @@ -378,9 +592,7 @@ def initialize_fixed_group_mixed_cga_work_id_state( """Register one physical fallback cluster and broadcast its fixed group coordinates.""" atomic_counter_state = work_id_state.atomic_counter_state if work_id_state.is_fallback_cluster: - broadcast_tensor = cute.make_tensor( - atomic_counter_state.broadcast_pointer, cute.make_layout((1,)) - ) + broadcast_tensor = cute.make_tensor(atomic_counter_state.broadcast_pointer, cute.make_layout((1,))) cluster_pipeline = atomic_counter_state.cluster_pipeline if atomic_counter_state.is_leader_cta: cluster_pipeline.producer_acquire(atomic_counter_state.producer_state) @@ -395,16 +607,11 @@ def initialize_fixed_group_mixed_cga_work_id_state( work_id_state.registration_counter_pointer, Int32(1), sem="relaxed", scope="gpu" ) fallback_ordinal = Int32( - cute.arch.shuffle_sync( - fallback_ordinal, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31 - ) + cute.arch.shuffle_sync(fallback_ordinal, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31) ) if lane_idx < Int32(atomic_counter_state.cluster_size): store_i32_to_peer_cluster_smem_async( - atomic_counter_state.broadcast_pointer, - fallback_ordinal, - full_barrier_pointer, - lane_idx, + atomic_counter_state.broadcast_pointer, fallback_ordinal, full_barrier_pointer, lane_idx ) mbarrier_arrive_expect_tx_on_peer(full_barrier_pointer, Int32(4), lane_idx) atomic_counter_state.producer_state.advance() @@ -412,16 +619,12 @@ def initialize_fixed_group_mixed_cga_work_id_state( cluster_pipeline.consumer_wait(atomic_counter_state.consumer_state) fallback_ordinal = broadcast_tensor[0] cute.arch.fence_acq_rel_cta() - cluster_pipeline.sync_object_empty.arrive( - atomic_counter_state.consumer_state.index, Int32(0) - ) + cluster_pipeline.sync_object_empty.arrive(atomic_counter_state.consumer_state.index, Int32(0)) atomic_counter_state.consumer_state.advance() fallback_group_idx = fallback_ordinal // Int32(work_id_state.split_factor) work_id_state.fallback_group_idx = fallback_group_idx - work_id_state.in_group_idx = fallback_ordinal - fallback_group_idx * Int32( - work_id_state.split_factor - ) + work_id_state.in_group_idx = fallback_ordinal - fallback_group_idx * Int32(work_id_state.split_factor) work_id_state.atomic_counter_state = atomic_counter_state return work_id_state @@ -440,17 +643,13 @@ def _claim_fixed_group_fallback_work_id( f"atomic_counter_index must be in [0, {atomic_counter_state.counter_count}), got {atomic_counter_index}." ) - broadcast_tensor = cute.make_tensor( - atomic_counter_state.broadcast_pointer, cute.make_layout((1,)) - ) + broadcast_tensor = cute.make_tensor(atomic_counter_state.broadcast_pointer, cute.make_layout((1,))) cluster_pipeline = atomic_counter_state.cluster_pipeline selected_counter_pointer = atomic_counter_state.counter_pointer + Int32(atomic_counter_index) if atomic_counter_state.is_leader_cta: cluster_pipeline.producer_acquire(atomic_counter_state.producer_state) - full_barrier_pointer = cluster_pipeline.sync_object_full.get_barrier( - atomic_counter_state.producer_state.index - ) + full_barrier_pointer = cluster_pipeline.sync_object_full.get_barrier(atomic_counter_state.producer_state.index) thread_idx, _, _ = cute.arch.thread_idx() lane_idx = thread_idx % Int32(32) group_base_offset = work_id_state.fallback_group_idx * Int32(work_id_state.split_factor) @@ -462,9 +661,7 @@ def _claim_fixed_group_fallback_work_id( while not all_members_consumed: observed_token = work_id_state.previous_token if lane_idx < Int32(work_id_state.split_factor): - observed_token = cute.arch.load( - group_token_pointer + lane_idx, Int64, sem="acquire", scope="gpu" - ) + observed_token = cute.arch.load(group_token_pointer + lane_idx, Int64, sem="acquire", scope="gpu") lane_is_ready = (lane_idx >= Int32(work_id_state.split_factor)) | ( observed_token == work_id_state.previous_token ) @@ -475,18 +672,12 @@ def _claim_fixed_group_fallback_work_id( claimed_work_id = Int32(0) if lane_idx == Int32(0): - claimed_work_id = cute.arch.atomic_add( - selected_counter_pointer, Int32(1), sem="relaxed", scope="gpu" - ) + claimed_work_id = cute.arch.atomic_add(selected_counter_pointer, Int32(1), sem="relaxed", scope="gpu") claimed_work_id = Int32( - cute.arch.shuffle_sync( - claimed_work_id, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31 - ) + cute.arch.shuffle_sync(claimed_work_id, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31) ) claimed_payload = claimed_work_id | (Int32(atomic_counter_index) << Int32(31)) - token = (Int64(work_id_state.next_generation) << Int64(32)) | ( - Int64(claimed_payload) & Int64(0xFFFFFFFF) - ) + token = (Int64(work_id_state.next_generation) << Int64(32)) | (Int64(claimed_payload) & Int64(0xFFFFFFFF)) if lane_idx == Int32(0): cute.arch.store(group_token_pointer, token, sem="relaxed", scope="gpu") work_id_state.previous_token = token @@ -497,36 +688,22 @@ def _claim_fixed_group_fallback_work_id( token_high = Int32(0) token_low = Int32(0) if lane_idx == Int32(0): - observed_token = cute.arch.load( - group_token_pointer, Int64, sem="relaxed", scope="gpu" - ) + observed_token = cute.arch.load(group_token_pointer, Int64, sem="relaxed", scope="gpu") token_high = Int32(observed_token >> Int64(32)) token_low = Int32(observed_token & Int64(0xFFFFFFFF)) - token_high = Int32( - cute.arch.shuffle_sync(token_high, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31) - ) - token_low = Int32( - cute.arch.shuffle_sync(token_low, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31) - ) + token_high = Int32(cute.arch.shuffle_sync(token_high, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31)) + token_low = Int32(cute.arch.shuffle_sync(token_low, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31)) token = (Int64(token_high) << Int64(32)) | (Int64(token_low) & Int64(0xFFFFFFFF)) if token == work_id_state.previous_token: nanosleep(500) claimed_payload = Int32(token & Int64(0xFFFFFFFF)) if lane_idx == Int32(0): - cute.arch.store( - group_token_pointer + work_id_state.in_group_idx, - token, - sem="relaxed", - scope="gpu", - ) + cute.arch.store(group_token_pointer + work_id_state.in_group_idx, token, sem="relaxed", scope="gpu") work_id_state.previous_token = token if lane_idx < Int32(atomic_counter_state.cluster_size): store_i32_to_peer_cluster_smem_async( - atomic_counter_state.broadcast_pointer, - claimed_payload, - full_barrier_pointer, - lane_idx, + atomic_counter_state.broadcast_pointer, claimed_payload, full_barrier_pointer, lane_idx ) mbarrier_arrive_expect_tx_on_peer(full_barrier_pointer, Int32(4), lane_idx) atomic_counter_state.producer_state.advance() @@ -555,9 +732,7 @@ def _claim_fixed_group_mixed_cga_work_id( work_id_state.atomic_counter_state = atomic_counter_state work_id_state.claimed_counter_index = Int32(atomic_counter_index) else: - linear_work_id, work_id_state = _claim_fixed_group_fallback_work_id( - work_id_state, atomic_counter_index - ) + linear_work_id, work_id_state = _claim_fixed_group_fallback_work_id(work_id_state, atomic_counter_index) return linear_work_id, work_id_state @@ -578,9 +753,7 @@ def _claim_cluster_launch_control_work_id( state_before_leader_query = work_id_state if work_id_state.is_leader_cta: work_id_state.cluster_pipeline.producer_acquire(work_id_state.producer_state) - response_barrier = work_id_state.cluster_pipeline.producer_get_barrier( - work_id_state.producer_state - ) + response_barrier = work_id_state.cluster_pipeline.producer_get_barrier(work_id_state.producer_state) with cute.arch.elect_one(): cute.arch.issue_clc_query(response_barrier, work_id_state.response_pointer) else: @@ -613,6 +786,68 @@ def _claim_cluster_launch_control_work_id( ) +@cute.jit +def _claim_clc_mixed_cga_work_id( + work_id_state: ClcMixedCgaWorkIdState, +) -> Tuple[GridWorkId, ClcMixedCgaWorkIdState]: + """Claim the next canceled cluster (first claim: own blockIdx) and return this CTA's grid coordinate.""" + use_bootstrap = work_id_state.response_pending + state_before_bootstrap = work_id_state + if use_bootstrap: + work_id_state.response_pending = Boolean(False) + else: + work_id_state = state_before_bootstrap + + state_before_query = work_id_state + if not use_bootstrap: + cluster_pipeline = work_id_state.cluster_pipeline + state_before_leader_query = work_id_state + if work_id_state.is_leader_cta: + cluster_pipeline.producer_acquire(work_id_state.producer_state) + full_barrier_pointer = cluster_pipeline.sync_object_full.get_barrier(work_id_state.producer_state.index) + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + if lane_idx < Int32(work_id_state.cluster_size): + # 16-byte CLC response. + mbarrier_arrive_expect_tx_on_peer(full_barrier_pointer, Int32(16), lane_idx) + with cute.arch.elect_one(): + cute.arch.issue_clc_query(full_barrier_pointer, work_id_state.response_pointer) + else: + work_id_state = state_before_leader_query + work_id_state.producer_state.advance() + + cluster_pipeline.consumer_wait(work_id_state.consumer_state) + (cluster_origin_m, cluster_origin_n, grid_l, response_is_valid) = cute.arch.clc_response( + work_id_state.response_pointer + ) + cute.arch.fence_acq_rel_cta() + cute.arch.fence_proxy("async.shared", space="cta") + # Cluster-scoped release orders this CTA's read of the response before the leader's next overwrite. + cute.arch.mbarrier_arrive( + cluster_pipeline.sync_object_empty.get_barrier(work_id_state.consumer_state.index), + Int32(0), + scope=_nvvm.MemScopeKind.CLUSTER, + ) + work_id_state.consumer_state.advance() + + work_id_state.grid_m = cluster_origin_m + work_id_state.cta_coord_in_cluster[0] + work_id_state.grid_n = cluster_origin_n + work_id_state.cta_coord_in_cluster[1] + work_id_state.grid_l = grid_l + work_id_state.response_is_valid = response_is_valid != Int32(0) + else: + work_id_state = state_before_query + + return ( + GridWorkId( + grid_m=work_id_state.grid_m, + grid_n=work_id_state.grid_n, + grid_l=work_id_state.grid_l, + is_valid=work_id_state.response_is_valid, + ), + work_id_state, + ) + + @cute.jit def claim_work_id(work_id_state, atomic_counter_index=0): """Claim the next work ID using the backend encoded by the state type.""" @@ -624,15 +859,20 @@ def claim_work_id(work_id_state, atomic_counter_index=0): return _claim_fixed_group_mixed_cga_work_id(work_id_state, atomic_counter_index) if cutlass.const_expr(isinstance(work_id_state, ClusterLaunchControlWorkIdState)): return _claim_cluster_launch_control_work_id(work_id_state) + if cutlass.const_expr(isinstance(work_id_state, ClcMixedCgaWorkIdState)): + return _claim_clc_mixed_cga_work_id(work_id_state) raise TypeError(f"Unsupported work-ID state: {type(work_id_state).__name__}.") __all__ = [ "AtomicCounterWorkIdState", + "ClcMixedCgaWorkIdState", "ClusterLaunchControlWorkIdState", "FixedGroupMixedCgaAtomicCounterWorkIdState", "GridStrideWorkIdState", "GridWorkId", + "broadcast_atomic_counter_ready", + "broadcast_atomic_counter_ready_or_claim", "claim_work_id", "initialize_fixed_group_mixed_cga_work_id_state", ] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/quant_def.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/quant_def.py index 82d5d0c7d11b..7ff5730d73c4 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/quant_def.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/quant_def.py @@ -56,10 +56,7 @@ def umma_kind(self) -> str: pair -- falls back to mxf8f6f4. Keeping the two in sync matters because we build the tiled MMA through that helper but emit the instruction ourselves. """ - both_fp4 = ( - self.weight_dtype is cutlass.Float4E2M1FN - and self.activation_dtype is cutlass.Float4E2M1FN - ) + both_fp4 = self.weight_dtype is cutlass.Float4E2M1FN and self.activation_dtype is cutlass.Float4E2M1FN if not both_fp4: return "mxf8f6f4" return "mxf4nvf4" if self.sf_vec_size == 16 else "mxf4" @@ -105,10 +102,7 @@ def needs_unpack_tma(self, architecture: str) -> bool: normalized_architecture = normalized_architecture[2:] if normalized_architecture not in ("100", "103", "107"): raise ValueError(f"Unsupported TCGen05 architecture {architecture!r}.") - return ( - normalized_architecture in ("100", "103") - and self.weight_dtype.width != self.activation_dtype.width - ) + return normalized_architecture in ("100", "103") and self.weight_dtype.width != self.activation_dtype.width @property def uses_global_scale(self) -> bool: @@ -145,23 +139,15 @@ def _instruction_descriptor_format_code(umma_kind: str, dtype: Type[cutlass.Nume try: return codes[dtype] except KeyError as error: - raise ValueError( - f"{dtype} has no instruction-descriptor format code under kind::{umma_kind}." - ) from error + raise ValueError(f"{dtype} has no instruction-descriptor format code under kind::{umma_kind}.") from error @dataclasses.dataclass(frozen=True) class CombineFormat: """Data and scale representation of one cross-rank FC2 payload.""" - _act_by_tag: ClassVar[Dict[str, type]] = { - "e2m1": cutlass.Float4E2M1FN, - "e4m3": cutlass.Float8E4M3FN, - } - _scale_by_tag: ClassVar[Dict[str, type]] = { - "bf16": cutlass.BFloat16, - "e8m0": cutlass.Float8E8M0FNU, - } + _act_by_tag: ClassVar[Dict[str, type]] = {"e2m1": cutlass.Float4E2M1FN, "e4m3": cutlass.Float8E4M3FN} + _scale_by_tag: ClassVar[Dict[str, type]] = {"bf16": cutlass.BFloat16, "e8m0": cutlass.Float8E8M0FNU} _rejection_reason: ClassVar[Dict[str, str]] = { "32e5m2xe8m0": ( "e5m2 costs the same 8.25 bits per element as e4m3 and trades a mantissa bit (6 dB of " @@ -200,9 +186,7 @@ def name(self) -> str: if not self.is_quantized: return "bf16" act_tag = next(tag for tag, dtype in self._act_by_tag.items() if dtype is self.act_dtype) - scale_tag = next( - tag for tag, dtype in self._scale_by_tag.items() if dtype is self.scale_dtype - ) + scale_tag = next(tag for tag, dtype in self._scale_by_tag.items() if dtype is self.scale_dtype) return f"{self.scale_block}{act_tag}x{scale_tag}" def __str__(self) -> str: @@ -217,9 +201,7 @@ def parse(cls, text: str) -> "CombineFormat": } token = text.strip().lower() if token in cls._rejection_reason: - raise ValueError( - f"Combine format {token!r} is deliberately unsupported: {cls._rejection_reason[token]}" - ) + raise ValueError(f"Combine format {token!r} is deliberately unsupported: {cls._rejection_reason[token]}") if token not in specs: raise ValueError(f"Invalid combine format {text!r}; expected one of {tuple(specs)}.") act_dtype, scale_dtype, scale_block = specs[token] @@ -252,15 +234,12 @@ def _verify_pinned_hardware_encoding() -> None: kind.scale_format_code, ) if derived != expected: - raise AssertionError( - f"QuantKind.{kind.name} derives {derived}, pinned encoding is {expected}." - ) + raise AssertionError(f"QuantKind.{kind.name} derives {derived}, pinned encoding is {expected}.") expected_instruction_k_2x = expected[2] * 2 instruction_k_2x = kind.instruction_k("2x") if instruction_k_2x != expected_instruction_k_2x: raise AssertionError( - f"QuantKind.{kind.name} derives 2x instruction K {instruction_k_2x}, " - f"expected {expected_instruction_k_2x}." + f"QuantKind.{kind.name} derives 2x instruction K {instruction_k_2x}, expected {expected_instruction_k_2x}." ) diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_kernels.py b/tensorrt_llm/_torch/moe/kernels/megamoe_kernels.py new file mode 100644 index 000000000000..abcc241122b6 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_kernels.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MegaMoE KernelClass entry points from the vendored ``cutedsl_megamoe`` package.""" + +from .cutedsl_megamoe.kernel_src.blackwell.inference.mega.block_scaled_swap_ab_mega_moe_kernel import ( + BlockScaledSwapAbMegaMoeKernel as BlackwellInferenceMegaMoE, +) +from .cutedsl_megamoe.kernel_src.rubin.inference.local_mega.block_scaled_swap_ab_local_mega_moe_kernel import ( + BlockScaledSwapAbLocalMegaMoeKernel as RubinInferenceLocalMegaMoE, +) +from .cutedsl_megamoe.kernel_src.rubin.inference.mega.block_scaled_swap_ab_mega_moe_kernel import ( + BlockScaledSwapAbMegaMoeKernel as RubinInferenceMegaMoE, +) +from .cutedsl_megamoe.kernel_src.rubin.inference.mega.block_scaled_swap_ab_mega_moe_kernel_gen_specialized import ( + BlockScaledSwapAbGenphaseMoeKernel as RubinInferenceGenphaseMegaMoE, +) + +__all__ = [ + "BlackwellInferenceMegaMoE", + "RubinInferenceGenphaseMegaMoE", + "RubinInferenceLocalMegaMoE", + "RubinInferenceMegaMoE", +] diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/VENDOR_STAMP b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/VENDOR_STAMP new file mode 100644 index 000000000000..f8e67c1f1393 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/VENDOR_STAMP @@ -0,0 +1,6 @@ +upstream: https://gitlab-master.nvidia.com/jintaop/cutedsl_eplb_scheduler_copy.git +locked_commit: 4dc8616031b18dd948bdf11b49afe0d86be15e22 +project: MegaMoE HALO-Q scheduler and SAMI runtime +original_authors: Jintao Peng ; Ethan Zhang +license: Apache-2.0 +managed_by: 3rdparty/vendor_sources.lock.yaml diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/__init__.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/__init__.py new file mode 100644 index 000000000000..1516768d758e --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/__init__.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""MegaMoE physical-slot scheduler public API.""" + +from .cuda_scheduler import ( + CudaPhysicalSlotScheduler, + CudaSchedulerConfig, + CudaSchedulerOutputs, + recommend_cuda_scheduler_ctas, +) +from .sami import ( + BoundHierarchicalLiveWeightArena, + HierarchicalCopyTicket, + HierarchicalLiveWeightArena, + HierarchicalLiveWeightArenaProvider, + HierarchicalSamiWeightBroadcast, + BoundLiveWeightArena, + BundleLayout, + Comm, + LivePlaneView, + LiveTerminalView, + LiveWeightArena, + LiveWeightArenaProvider, + hierarchy_group_sizes, +) + +__all__ = [ + "BoundHierarchicalLiveWeightArena", + "CudaPhysicalSlotScheduler", + "CudaSchedulerConfig", + "CudaSchedulerOutputs", + "recommend_cuda_scheduler_ctas", + "BoundLiveWeightArena", + "BundleLayout", + "Comm", + "LivePlaneView", + "HierarchicalCopyTicket", + "HierarchicalLiveWeightArena", + "HierarchicalLiveWeightArenaProvider", + "HierarchicalSamiWeightBroadcast", + "LiveTerminalView", + "LiveWeightArena", + "LiveWeightArenaProvider", + "hierarchy_group_sizes", +] diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/__init__.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/__init__.py new file mode 100644 index 000000000000..90410ba4fbc1 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/__init__.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Pure-CUDA physical-slot scheduler. + +HALO-M hierarchical placement followed by HALO-Q quota repair, behind one +four-tensor placement ABI. Copy routing belongs to the in-switch copy layer. The CPU oracle that states the same placement lives in +``research/halo_q_oracle.py`` and is not imported from here. +""" + +from .runtime import ( + CudaPhysicalSlotScheduler, + CudaSchedulerConfig, + CudaSchedulerOutputs, + recommend_cuda_scheduler_ctas, +) + +__all__ = [ + "CudaPhysicalSlotScheduler", + "CudaSchedulerConfig", + "CudaSchedulerOutputs", + "recommend_cuda_scheduler_ctas", +] diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/multirank_dist.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/multirank_dist.py new file mode 100644 index 000000000000..7074513cb639 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/multirank_dist.py @@ -0,0 +1,125 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CUDA fabric allocation shared across expert-parallel processes.""" + +from __future__ import annotations + +import torch +from cuda.bindings import driver as cuda + + +def _check(result: object) -> object: + if isinstance(result, tuple): + error, *values = result + else: + error, values = result, [] + if int(error) != int(cuda.CUresult.CUDA_SUCCESS): + _, name = cuda.cuGetErrorName(error) + _, description = cuda.cuGetErrorString(error) + if isinstance(name, bytes): + name = name.decode() + if isinstance(description, bytes): + description = description.decode() + raise RuntimeError(f"CUDA driver call failed: {name}: {description}") + if not values: + return None + return values[0] if len(values) == 1 else tuple(values) + + +class FabricSymmetricBuffer: + """One CUDA-fabric allocation mapped read/write into every EP process.""" + + def __init__(self, nbytes: int, device: int): + self.device = device + properties = cuda.CUmemAllocationProp() + properties.type = cuda.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED + properties.location.type = cuda.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + properties.location.id = device + properties.requestedHandleTypes = ( + cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC + ) + self._properties = properties + self.granularity = int( + _check( + cuda.cuMemGetAllocationGranularity( + properties, + cuda.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED, + ) + ) + ) + self.size = (nbytes + self.granularity - 1) // self.granularity + self.size *= self.granularity + self._mappings: list[int] = [] + self._reservations: list[int] = [] + self._imported: list[object] = [] + self._closed = False + self._handle = None + self.ptr = 0 + self.shareable = b"" + try: + self._handle = _check(cuda.cuMemCreate(self.size, properties, 0)) + self.ptr = self._map(self._handle) + _check(cuda.cuMemsetD8(cuda.CUdeviceptr(self.ptr), 0, self.size)) + exported = _check( + cuda.cuMemExportToShareableHandle( + self._handle, + cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC, + 0, + ) + ) + self.shareable = bytes(exported.data) + except BaseException: + self.close() + raise + + def _map(self, handle: object) -> int: + address = int( + _check(cuda.cuMemAddressReserve(self.size, self.granularity, 0, 0)) + ) + self._reservations.append(address) + _check(cuda.cuMemMap(address, self.size, 0, handle, 0)) + self._mappings.append(address) + access = cuda.CUmemAccessDesc() + access.location.type = cuda.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + access.location.id = self.device + access.flags = cuda.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE + _check(cuda.cuMemSetAccess(address, self.size, [access], 1)) + return int(address) + + def import_peer(self, shareable: bytes) -> int: + handle = _check( + cuda.cuMemImportFromShareableHandle( + shareable, + cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC, + ) + ) + self._imported.append(handle) + return self._map(handle) + + def close(self) -> None: + """Release every mapping and fabric handle after device work is drained.""" + if self._closed: + return + with torch.cuda.device(self.device): + while self._mappings: + address = self._mappings[-1] + _check(cuda.cuMemUnmap(address, self.size)) + self._mappings.pop() + while self._reservations: + address = self._reservations[-1] + _check(cuda.cuMemAddressFree(address, self.size)) + self._reservations.pop() + while self._imported: + handle = self._imported[-1] + _check(cuda.cuMemRelease(handle)) + self._imported.pop() + if self._handle is not None: + _check(cuda.cuMemRelease(self._handle)) + self._handle = None + self.ptr = 0 + self.shareable = b"" + self._closed = True + + +__all__ = ["FabricSymmetricBuffer"] diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py new file mode 100644 index 000000000000..b88af7be64ba --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py @@ -0,0 +1,607 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Host runtime for the pure-CUDA HALO-Q scheduler.""" + +from __future__ import annotations + +from dataclasses import dataclass +import threading + +import torch + +from ..geometry import MAX_EP, MIN_EP, hierarchy_group_sizes + + +MAX_EXPERTS = 384 +MAX_CTAS = 128 +MAX_BROADCASTS = MAX_EXPERTS +SPIN_CYCLES = 8_000_000_000 + +def _plan_channel_word_count(*, abi_version: int, channel_words: int | None, + route_features: int, ep: int, helpers: int) -> int: + """Validate an explicitly negotiated private publication allocation.""" + if type(abi_version) is not int or type(route_features) is not int: + raise ValueError("plan channel ABI and features must be exact integers") + if abi_version != 6: + raise ValueError("CUDA scheduler requires PlanChannel ABI6") + if route_features not in (0, 1): + raise ValueError("plan channel route features must be 0 or 1") + if type(ep) is not int or not MIN_EP <= ep <= MAX_EP or ep % 2: + raise ValueError("plan channel requires even EP in [2,32]") + if type(helpers) is not int or helpers <= 0: + raise ValueError("helper count must be a positive exact int") + stride = (helpers + 3) // 4 * 4 + expected = 4 + 7 * stride + if expected > 2**31 - 1: + raise ValueError("plan channel word count does not fit its int32 ABI") + if channel_words is None: + channel_words = expected + if type(channel_words) is not int or channel_words != expected: + raise ValueError("plan channel allocation does not match its explicit ABI") + return channel_words + + +def recommend_cuda_scheduler_ctas( + ep_size: int, + max_tokens_per_rank: int, + topk: int = 6, + threads: int = 512, +) -> int: + """Return the bounded pure-CUDA CTA policy without importing CuTe.""" + + for name, value in { + "ep_size": ep_size, + "max_tokens_per_rank": max_tokens_per_rank, + "topk": topk, + "threads": threads, + }.items(): + if type(value) is not int or value <= 0: + raise ValueError(f"{name} must be a positive exact int") + if threads != 512: + raise ValueError("the pure-CUDA scheduler currently requires threads=512") + routes = max_tokens_per_rank * topk + if ep_size == 8: + # Bound scheduler occupancy so independent same-stream work retains + # launch capacity. + workers = min(31, max(1, (routes + 3 * threads - 1) // (3 * threads))) + return 1 + workers + target_rounds = max(3, (96 + (ep_size + 6) // 2) // (ep_size + 6)) + workers = max( + 1, (max_tokens_per_rank * topk) // (threads * target_rounds) + ) + workers = min(workers, MAX_CTAS - 1) + rounds = (routes + workers * threads - 1) // (workers * threads) + workers = (routes + rounds * threads - 1) // (rounds * threads) + return 1 + workers + + +def symmetric_buffer_ints(ep_size: int, expert_count: int) -> int: + # Two reset-free arrival generations plus a local epoch and two count slabs. + return 64 + 2 * ep_size * expert_count + + +@dataclass(frozen=True) +class CudaSchedulerConfig: + ep_size: int + logical_expert_count: int + extra_slots_per_rank: int + max_tokens_per_rank: int + topk: int = 6 + local_rank: int = 0 + threads: int = 512 + ctas: int | None = None + enable_pdl: bool = True + + def __post_init__(self) -> None: + if self.ctas is None: + object.__setattr__( + self, + "ctas", + recommend_cuda_scheduler_ctas( + self.ep_size, + self.max_tokens_per_rank, + self.topk, + self.threads, + ), + ) + self.validate() + + @property + def home_experts_per_rank(self) -> int: + return self.logical_expert_count // self.ep_size + + @property + def local_physical_slot_count(self) -> int: + return self.home_experts_per_rank + self.extra_slots_per_rank + + @property + def route_count(self) -> int: + return self.max_tokens_per_rank * self.topk + + @property + def world_size(self) -> int: + return self.ep_size + + @property + def global_expert_count(self) -> int: + return self.logical_expert_count + + @property + def helper_count(self) -> int: + return self.extra_slots_per_rank + + @property + def group_sizes(self) -> tuple[int, ...]: + return hierarchy_group_sizes(self.ep_size) + + def validate(self) -> None: + values = { + "ep_size": self.ep_size, + "logical_expert_count": self.logical_expert_count, + "extra_slots_per_rank": self.extra_slots_per_rank, + "max_tokens_per_rank": self.max_tokens_per_rank, + "topk": self.topk, + "local_rank": self.local_rank, + "threads": self.threads, + "ctas": self.ctas, + } + if any(type(value) is not int for value in values.values()): + raise ValueError("all CUDA scheduler dimensions must be exact ints") + if not MIN_EP <= self.ep_size <= MAX_EP or self.ep_size % 2: + raise ValueError( + f"pure-CUDA scheduler supports even EP in [{MIN_EP},{MAX_EP}]") + if not 0 <= self.local_rank < self.ep_size: + raise ValueError("local_rank must be in [0,ep_size)") + if not 1 <= self.logical_expert_count <= MAX_EXPERTS: + raise ValueError(f"logical_expert_count must be in [1,{MAX_EXPERTS}]") + if self.logical_expert_count % self.ep_size: + raise ValueError("logical_expert_count must be EP-divisible") + if self.extra_slots_per_rank <= 0: + raise ValueError("extra_slots_per_rank must be positive") + if self.ep_size * self.local_physical_slot_count > 2**31 - 1: + raise ValueError("global physical slot addresses must fit int32") + if self.max_tokens_per_rank <= 0 or self.topk <= 0: + raise ValueError("token and topk dimensions must be positive") + if self.ep_size * self.route_count > 2**31 - 1: + raise ValueError("global route count must fit int32 histogram totals") + if self.threads != 512: + raise ValueError("the pure-CUDA scheduler currently requires threads=512") + if not 1 <= self.ctas <= MAX_CTAS: + raise ValueError(f"ctas must be in [1,{MAX_CTAS}]") + if type(self.enable_pdl) is not bool: + raise ValueError("enable_pdl must be bool") + + +class CudaSchedulerOutputs: + """The exact four-tensor hierarchical scheduler ABI.""" + + __slots__ = ( + "physical_slot_ids", + "hot_expert_ids", + "hot_expert_group_level", + "hot_expert_source_ranks", + ) + + def __init__(self, config: CudaSchedulerConfig, device: torch.device): + self.physical_slot_ids = torch.empty( + (config.max_tokens_per_rank, config.topk), + dtype=torch.int32, + device=device, + ) + shape = (config.extra_slots_per_rank,) + self.hot_expert_ids = torch.empty(shape, dtype=torch.int32, device=device) + self.hot_expert_group_level = torch.empty( + shape, dtype=torch.int32, device=device + ) + self.hot_expert_source_ranks = torch.empty( + shape, dtype=torch.int32, device=device + ) + + +def _stream_handle(stream: object) -> int: + if isinstance(stream, bool): + raise TypeError("stream must be a CUDA stream handle") + handle = getattr(stream, "cuda_stream", stream) + try: + value = int(handle) + except (TypeError, ValueError) as error: + raise TypeError("stream must be a CUDA stream handle") from error + if value < 0: + raise ValueError("stream handle cannot be negative") + return value + + +class CudaPhysicalSlotScheduler: + """Preallocated pure-CUDA scheduler with no CuTe JIT executor.""" + + def __init__(self, config: CudaSchedulerConfig, device: str | torch.device = "cuda"): + config.validate() + scheduler_device = torch.device(device) + if scheduler_device.type != "cuda": + raise ValueError("CudaPhysicalSlotScheduler requires a CUDA device") + if scheduler_device.index is None: + scheduler_device = torch.device("cuda", torch.cuda.current_device()) + properties = torch.cuda.get_device_properties(scheduler_device) + if config.ctas > properties.multi_processor_count: + raise ValueError("ctas cannot exceed the CUDA device SM count") + + self.cfg = config + self.device = scheduler_device + self.outputs = CudaSchedulerOutputs(config, scheduler_device) + self.out = self.outputs + self.logical_expert_ids = torch.full( + (config.max_tokens_per_rank, config.topk), + -1, + dtype=torch.int32, + device=scheduler_device, + ) + self.status = torch.zeros(8, dtype=torch.int32, device=scheduler_device) + # Everything the kernel takes between `routes` and `stream` is fixed for + # the endpoint's lifetime. Cache it so a submission does not rebuild 13 + # data_ptr() calls and ~11 config lookups every generation. + self.partial = torch.zeros( + config.ctas * config.logical_expert_count, + dtype=torch.int32, + device=scheduler_device, + ) + warps = config.threads // 32 + bins = config.logical_expert_count + 1 + route_aux_ints = config.ctas * warps * bins + if config.ctas > 1: + route_aux_ints = max( + route_aux_ints, + 2 + (config.ctas - 1) * bins + config.route_count, + ) + self.route_aux = torch.zeros( + route_aux_ints, + dtype=torch.int32, + device=scheduler_device, + ) + # The main barrier uses [arrival, epoch]; word 2 publishes route readiness. + self.grid_sync = torch.zeros(4, dtype=torch.int32, device=scheduler_device) + # Each expert is broadcast at most once. The private plan's fixed + # field stride is bounded by E, independently of requested slot capacity. + self.max_broadcasts = MAX_BROADCASTS + self.plan_workspace = torch.full( + (2 + 6 * self.max_broadcasts,), + -1, + dtype=torch.int32, + device=scheduler_device, + ) + self.route_prefix = torch.zeros( + self.max_broadcasts * (config.ep_size + 1), + dtype=torch.int32, + device=scheduler_device, + ) + self.sym = torch.zeros( + symmetric_buffer_ints( + config.ep_size, config.logical_expert_count + ), + dtype=torch.int32, + device=scheduler_device, + ) + self.peer_base = torch.full( + (config.ep_size,), + int(self.sym.data_ptr()), + dtype=torch.int64, + device=scheduler_device, + ) + self._connected = config.ep_size == 1 + self._plan_channel_lock = threading.Lock() + self._plan_channel_ptr = 0 + self._plan_channel_abi_version = 0 + self._plan_channel_words = 0 + self._plan_channel_route_features = 0 + self._plan_channel_bound = False + self._gpu_direct_bound = False + self._plan_channel_pending = False + self._plan_channel_stream_handle: int | None = None + self._plan_channel_torch_stream: torch.cuda.Stream | None = None + self._launch_stream_handle: int | None = None + + def _staging_stream(self) -> torch.cuda.Stream: + current = torch.cuda.current_stream(self.device) + bound = self._plan_channel_torch_stream + if bound is not None and _stream_handle(current) != self._plan_channel_stream_handle: + bound.wait_stream(current) + return bound + return current + + def stage(self, logical_expert_ids: torch.Tensor, *, validate_values: bool = False) -> None: + expected = (self.cfg.max_tokens_per_rank, self.cfg.topk) + if not isinstance(logical_expert_ids, torch.Tensor): + raise TypeError("logical_expert_ids must be a torch.Tensor") + if tuple(logical_expert_ids.shape) != expected: + raise ValueError(f"logical_expert_ids must have shape {expected}") + if logical_expert_ids.dtype != torch.int32 or not logical_expert_ids.is_contiguous(): + raise ValueError("logical_expert_ids must be contiguous torch.int32") + if not logical_expert_ids.is_cuda or logical_expert_ids.device != self.device: + raise ValueError("logical_expert_ids must be on the scheduler CUDA device") + if validate_values: + invalid = (logical_expert_ids < -1) | ( + logical_expert_ids >= self.cfg.logical_expert_count + ) + if bool(invalid.any().item()): + raise ValueError("logical expert IDs must be -1 or in [0,E)") + with torch.cuda.device(self.device): + stream = self._staging_stream() + with torch.cuda.stream(stream): + # No status clear: the kernel writes the failing generation into + # each flag rather than a bare 1, so a stale flag is + # distinguishable from a fresh one and clearing is only needed to + # acknowledge a fault (`reset_status()`). + self.logical_expert_ids.copy_(logical_expert_ids) + logical_expert_ids.record_stream(stream) + + def connect_peer_bases(self, peer_bases: object) -> None: + if self._connected: + raise RuntimeError("peer base binding is single-assignment") + + bases = tuple(peer_bases) # type: ignore[arg-type] + if len(bases) != self.cfg.ep_size or any( + type(pointer) is not int or pointer <= 0 for pointer in bases + ): + raise ValueError("peer_bases must contain EP positive exact addresses") + self.peer_base.copy_( + torch.tensor(bases, dtype=torch.int64, device=self.device) + ) + torch.cuda.synchronize(self.device) + self._connected = True + + def bind_plan_channel( + self, + device_ptr: int, + stream: object, + *, + abi_version: int, + channel_words: int | None = None, + route_features: int = 0, + ) -> None: + with self._plan_channel_lock: + words = _plan_channel_word_count( + abi_version=abi_version, channel_words=channel_words, + route_features=route_features, ep=self.cfg.ep_size, + helpers=self.cfg.extra_slots_per_rank, + ) + if self._plan_channel_bound: + raise RuntimeError("plan channel binding is single-assignment") + if type(device_ptr) is not int or device_ptr <= 0 or device_ptr % 16: + raise ValueError("device_ptr must be a positive 16-byte aligned address") + handle = _stream_handle(stream) + with torch.cuda.device(self.device): + self._plan_channel_torch_stream = torch.cuda.ExternalStream( + handle, device=self.device + ) + self._plan_channel_ptr = device_ptr + self._plan_channel_abi_version = abi_version + self._plan_channel_words = words + self._plan_channel_route_features = route_features + self._plan_channel_stream_handle = handle + self._plan_channel_bound = True + + def bind_gpu_direct(self, stream: object) -> None: + """Bind a copy consumer to the existing outputs/workspace, without publication. + + The consumer must enqueue on this static stream, then call + ``release_plan_channel`` before the next scheduler submission. + No additional plan allocation or GPU/host transfer is performed. + """ + with self._plan_channel_lock: + if self._plan_channel_bound: + raise RuntimeError("plan channel binding is single-assignment") + handle = _stream_handle(stream) + if self._launch_stream_handle is not None and handle != self._launch_stream_handle: + raise RuntimeError("GPU-direct binding requires the existing static stream") + with torch.cuda.device(self.device): + self._plan_channel_torch_stream = torch.cuda.ExternalStream(handle, device=self.device) + self._plan_channel_ptr = 0 + self._plan_channel_abi_version = 0 + self._plan_channel_words = 0 + self._plan_channel_route_features = 0 + self._plan_channel_stream_handle = handle + self._plan_channel_bound = True + self._gpu_direct_bound = True + + def release_plan_channel(self) -> None: + """Allow another publication after its copy has been enqueued. + + GPU-direct readers and scheduler writers use the same static stream, + so stream order protects the single plan allocation without a host + completion wait. This does not release the weight bank's consumer lease. + """ + with self._plan_channel_lock: + if not self._plan_channel_bound or not self._plan_channel_pending: + raise RuntimeError("plan channel has no pending publication") + self._plan_channel_pending = False + + def launch(self, stream: object | None = None) -> CudaSchedulerOutputs: + if not self._connected: + raise RuntimeError("EP peer bases are not connected") + with self._plan_channel_lock: + if self._plan_channel_bound: + if self._plan_channel_pending: + raise RuntimeError("plan publication is still pending") + handle = self._plan_channel_stream_handle + if stream is not None and _stream_handle(stream) != handle: + raise RuntimeError("bound plan channel requires its static stream") + else: + handle = _stream_handle( + torch.cuda.current_stream(self.device) if stream is None else stream + ) + if self._launch_stream_handle is None: + self._launch_stream_handle = handle + elif handle != self._launch_stream_handle: + raise RuntimeError( + "CUDA scheduler uses one static stream because its workspace " + "and reusable barrier epochs are shared" + ) + self._launch_op( + self.logical_expert_ids, + int(handle), + valid_route_count=self.cfg.route_count, + ) + if self._plan_channel_bound: + self._plan_channel_pending = True + return self.outputs + + def _launch_op( + self, + routes: torch.Tensor, + stream_handle: int, + *, + valid_route_count: int, + ) -> None: + stream = self._plan_channel_torch_stream + if stream is None or _stream_handle(stream) != stream_handle: + stream = torch.cuda.ExternalStream(stream_handle, device=self.device) + with torch.cuda.device(self.device), torch.cuda.stream(stream): + torch.ops.trtllm.moe_rebalance_halo_q( + routes, + self.outputs.physical_slot_ids, + self.outputs.hot_expert_ids, + self.outputs.hot_expert_group_level, + self.outputs.hot_expert_source_ranks, + self.peer_base, + self.status, + self.partial, + self.route_aux, + self.grid_sync, + self.plan_workspace, + self.route_prefix, + self._plan_channel_ptr, + self.cfg.ep_size, + self.cfg.logical_expert_count, + self.cfg.extra_slots_per_rank, + self.cfg.local_rank, + self.cfg.route_count, + self.cfg.ctas, + self.cfg.threads, + self.cfg.enable_pdl, + SPIN_CYCLES, + self._plan_channel_abi_version, + self._plan_channel_words, + self._plan_channel_route_features, + valid_route_count, + ) + + def submit( + self, routes: torch.Tensor, stream_handle: int, *, + valid_tokens: int | None = None, + ) -> CudaSchedulerOutputs: + """Read a caller-owned CUDA int32[T,K] tensor without input staging. + + T may be zero and may differ across ranks, up to max_tokens_per_rank. + By default all T rows are valid; valid_tokens can select a shorter + prefix of an existing allocation. The kernel never reads outside that + prefix and writes -1 to the unused output tail. All ranks, including + empty ranks, must submit each generation in the same collective order. + + The caller must order input production before this static stream and + keep routes alive and unmodified until the scheduler completes. This + includes allocator lifetime protection for temporary cross-stream + inputs. Metadata checks below do not launch GPU work or synchronize; + invalid expert values are reported by the kernel through check_status(). + + Graph capture fixes both the input pointer and valid_tokens. For replay + with changing pointers/lengths use eager submit; stage()+launch() remains + available for a fixed-capacity graph input buffer. + """ + with self._plan_channel_lock: + if self._plan_channel_bound: + if self._plan_channel_pending: + raise RuntimeError("plan publication is still pending") + if stream_handle != self._plan_channel_stream_handle: + raise RuntimeError("bound plan channel requires its static stream") + if self._launch_stream_handle is None: + self._launch_stream_handle = stream_handle + elif stream_handle != self._launch_stream_handle: + raise RuntimeError( + "CUDA scheduler uses one static stream because its workspace " + "and reusable barrier epochs are shared" + ) + self._submit_routes(routes, stream_handle, valid_tokens=valid_tokens) + if self._plan_channel_bound: + self._plan_channel_pending = True + return self.outputs + + def _submit_routes( + self, routes: torch.Tensor, stream_handle: int, *, + valid_tokens: int | None = None, + ) -> None: + if not isinstance(routes, torch.Tensor): + raise TypeError("routes must be a torch.Tensor") + if (routes.ndim != 2 or routes.shape[1] != self.cfg.topk + or routes.shape[0] > self.cfg.max_tokens_per_rank): + raise ValueError("routes must have shape [T,topk] with 0 <= T <= max_tokens_per_rank") + if routes.dtype != torch.int32 or not routes.is_contiguous(): + raise ValueError("routes must be contiguous torch.int32") + if not routes.is_cuda or routes.device != self.device: + raise ValueError("routes must be on the scheduler CUDA device") + if valid_tokens is None: + valid_tokens = routes.shape[0] + if type(valid_tokens) is not int or not 0 <= valid_tokens <= routes.shape[0]: + raise ValueError("valid_tokens must be an int in [0,T]") + self._launch_op( + routes, + stream_handle, + valid_route_count=valid_tokens * self.cfg.topk, + ) + + def run( + self, + logical_expert_ids: torch.Tensor, + stream: object | None = None, + *, + validate_values: bool = False, + ) -> CudaSchedulerOutputs: + if stream is not None and not self._plan_channel_bound: + with torch.cuda.device(self.device): + external = torch.cuda.ExternalStream( + _stream_handle(stream), device=self.device + ) + with torch.cuda.stream(external): + self.stage(logical_expert_ids, validate_values=validate_values) + else: + self.stage(logical_expert_ids, validate_values=validate_values) + return self.launch(stream) + + def check_status(self) -> None: + """Raise if any generation faulted since the last `reset_status()`. + + Each flag holds the unsigned modular epoch that wrote it, not a bare 1. + Zero means never set; an error at wrapped epoch zero stores UINT32_MAX. + That lets the per-generation status clear stay out of the submission path. + """ + torch.cuda.synchronize(self.device) + flags = [ + flag & 0xFFFFFFFF for flag in self.status[:6].tolist() + ] # one D2H, not six + labels = ( + "cross-rank histogram rendezvous timeout", + "CTA grid barrier timeout", + "plan publication timeout", + "invalid logical expert ID", + "HALO helper coloring failure", + "HALO quota/host invariant failure", + ) + failures = [ + f"{label} (generation {flag})" + for label, flag in zip(labels, flags) + if flag + ] + if failures: + raise RuntimeError("; ".join(failures)) + + def reset_status(self) -> None: + """Acknowledge and clear faults. Not needed per generation.""" + self.status.zero_() + + +__all__ = [ + "CudaPhysicalSlotScheduler", + "CudaSchedulerConfig", + "CudaSchedulerOutputs", + "recommend_cuda_scheduler_ctas", + "symmetric_buffer_ints", +] diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/geometry.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/geometry.py new file mode 100644 index 000000000000..f1fb15739e9d --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/geometry.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The R4 rank hierarchy shared by the CUDA scheduler and the SAMI copy.""" + +from __future__ import annotations + +__all__ = ["MAX_EP", "MIN_EP", "MIN_GROUP_SIZE", "hierarchy_group_sizes"] + +MIN_EP = 2 +MAX_EP = 32 +MIN_GROUP_SIZE = 4 + + +def hierarchy_group_sizes(world: int) -> tuple[int, ...]: + """Return HALO-M's R4 aligned hierarchy, coarsest to finest. + + Level 0 is the whole EP world. Later levels halve the group size only + while the world divides exactly and every group still holds at least + ``MIN_GROUP_SIZE`` ranks. ``halo_q_scheduler.cu::hierarchy_depth`` is the + device-side statement of the same ladder; the two must agree for every + world this package accepts, so there is deliberately one host definition + rather than one per subpackage. + """ + + if (type(world) is not int or world < MIN_EP or world > MAX_EP + or world % 2): + raise ValueError( + f"hierarchical SAMI supports even EP in [{MIN_EP},{MAX_EP}]") + result = [world] + groups = 2 + while world % groups == 0 and world // groups >= MIN_GROUP_SIZE: + result.append(world // groups) + groups *= 2 + return tuple(result) diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/__init__.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/__init__.py new file mode 100644 index 000000000000..851fe65a18f0 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/__init__.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Production SAMI resources for an external single-bank live arena.""" + +from .arena import ( + BoundLiveWeightArena, + LivePlaneView, + LiveTerminalView, + LiveWeightArena, + LiveWeightArenaProvider, +) +from ._util import Comm +from .geometry import BundleLayout, Plane, PlaneViewLayout +from .hierarchical import ( + BoundHierarchicalLiveWeightArena, + HierarchicalCopyTicket, + HierarchicalLiveWeightArena, + HierarchicalLiveWeightArenaProvider, + HierarchicalSamiWeightBroadcast, + hierarchy_group_sizes, +) + +__all__ = [ + "BoundHierarchicalLiveWeightArena", + "BoundLiveWeightArena", + "BundleLayout", + "Comm", + "HierarchicalCopyTicket", + "HierarchicalLiveWeightArena", + "HierarchicalLiveWeightArenaProvider", + "HierarchicalSamiWeightBroadcast", + "LivePlaneView", + "LiveTerminalView", + "LiveWeightArena", + "LiveWeightArenaProvider", + "Plane", + "PlaneViewLayout", + "hierarchy_group_sizes", +] diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/_util.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/_util.py new file mode 100644 index 000000000000..a97ad6b6692a --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/_util.py @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The driver-call wrapper and setup collectives production SAMI needs. + +Two small helpers rather than two small modules. They are unrelated to each +other, but both are leaf utilities with a single consumer set, and neither is +part of the package's public API. + +``check_cuda`` deliberately does not touch torch. That separation is not +load-bearing for import cost -- ``sami/__init__.py`` already imports ``Comm`` +and ``hierarchical``, so torch is on the path for any ``sami.*`` import -- but +it does keep the driver wrapper usable from a torch-free context if the +package layout ever changes. +""" + +from __future__ import annotations + +from cuda.bindings import driver as cuda +import torch.distributed as dist + +__all__ = ["Comm", "check_cuda"] + + +def check_cuda(result: object, operation: str = "CUDA") -> object: + """Unpack a cuda-python driver result and raise on failure.""" + if isinstance(result, tuple): + error, *values = result + else: + error, values = result, [] + if int(error) != int(cuda.CUresult.CUDA_SUCCESS): + _, name = cuda.cuGetErrorName(error) + _, description = cuda.cuGetErrorString(error) + if isinstance(name, bytes): + name = name.decode() + if isinstance(description, bytes): + description = description.decode() + raise RuntimeError( + f"{operation} failed: {name}: {description} (code {int(error)})" + ) + if not values: + return None + return values[0] if len(values) == 1 else tuple(values) + + +class Comm: + """The rank identity and setup collectives required by SAMI.""" + + def __init__(self) -> None: + if not dist.is_initialized(): + raise RuntimeError("torch.distributed must be initialized before SAMI") + self.rank = dist.get_rank() + self.world = dist.get_world_size() + + def barrier(self) -> None: + dist.barrier() + + def bcast(self, payload: object, root: int) -> object: + values = [payload] + dist.broadcast_object_list(values, src=root) + return values[0] + + def allgather(self, value: object) -> list[object]: + values: list[object] = [None] * self.world + dist.all_gather_object(values, value) + return values diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/arena.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/arena.py new file mode 100644 index 000000000000..20aef7bbaba3 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/arena.py @@ -0,0 +1,419 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""External single-bank MegaMoE weight arena consumed directly by SAMI. + +The production arena is plane-major. Each of the seven tensors has local +physical slots ``[0, H)`` for canonical experts and ``[H, H + S)`` for live +helpers. A plane exposes every rank's unicast mapping plus the multicast +mapping of the same backing allocation. This is the shape naturally exposed +by TEKit ``McastGPUBuffer.get_uc_buffer``/``get_mc_buffer``; callers without +that class can provide equivalent typed CUDA tensor views. + +This module performs no allocation and deliberately does not import TEKit. +It validates geometry, pointer and object identity once at bind time and +resolves the pointer tables the copy backend uses from then on; identity is +not re-checked per submit. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Protocol, runtime_checkable + +from .geometry import BundleLayout + + +def _exact_positive_int(name: str, value: object) -> int: + if type(value) is not int or value <= 0: + raise ValueError(f"{name} must be a positive exact int") + return value + + +def _pointer(value: object, *, name: str) -> int: + accessor = getattr(value, "data_ptr", None) + pointer = accessor() if callable(accessor) else accessor + if type(pointer) is not int or pointer <= 0: + raise TypeError(f"{name} must expose a positive integer data_ptr") + return pointer + + +def _shape(value: object, *, name: str) -> tuple[int, ...]: + try: + result = tuple(int(extent) for extent in value.shape) # type: ignore[attr-defined] + except (AttributeError, TypeError, ValueError) as exc: + raise TypeError(f"{name} must expose an integer shape") from exc + if not result or any(extent <= 0 for extent in result): + raise ValueError(f"{name} must have a non-empty positive shape") + return result + + +def _stride(value: object, *, name: str) -> tuple[int, ...]: + accessor = getattr(value, "stride", None) + raw = accessor() if callable(accessor) else accessor + try: + result = tuple(int(step) for step in raw) + except (TypeError, ValueError) as exc: + raise TypeError(f"{name} must expose an integer stride") from exc + if any(step <= 0 for step in result): + raise ValueError(f"{name} must have positive strides") + return result + + +def _dtype(value: object) -> str: + return str(getattr(value, "dtype", None)).lower().removeprefix("torch.") + + +def _element_size(value: object, *, name: str) -> int: + accessor = getattr(value, "element_size", None) + raw = accessor() if callable(accessor) else accessor + if type(raw) is not int or raw <= 0: + raise TypeError(f"{name} must expose a positive integer element_size") + return raw + + +def _device_index(value: object, *, name: str) -> int: + if getattr(value, "is_cuda", True) is not True: + raise ValueError(f"{name} must reside on CUDA memory") + device = getattr(value, "device", None) + text = str(device).lower() + if not text.startswith("cuda"): + raise ValueError(f"{name} must reside on CUDA memory") + index = getattr(device, "index", None) + if type(index) is not int: + if ":" not in text: + raise ValueError(f"{name} must expose an explicit CUDA device ordinal") + try: + index = int(text.split(":", 1)[1]) + except ValueError as exc: + raise ValueError(f"{name} exposes an invalid CUDA device") from exc + if index < 0: + raise ValueError(f"{name} exposes an invalid CUDA device") + return index + + +@dataclass(frozen=True) +class LivePlaneView: + """One live plane's rank-ordered UC views and its MC alias. + + ``aliases_same_backing`` is an explicit provider attestation. Pointer + arithmetic alone cannot prove a CUDA UC/MC alias relationship, so an + un-attested collection of unrelated tensors is rejected. + """ + + name: str + uc_views: tuple[object, ...] = field(repr=False, compare=False) + mc_view: object = field(repr=False, compare=False) + aliases_same_backing: bool = False + backing_owner: object | None = field(default=None, repr=False, compare=False) + + +@dataclass(frozen=True) +class LiveTerminalView: + """Local ``uint64[EP]`` terminals and the matching MC alias.""" + + local_view: object = field(repr=False, compare=False) + mc_view: object = field(repr=False, compare=False) + aliases_same_backing: bool = False + backing_owner: object | None = field(default=None, repr=False, compare=False) + + +@dataclass(frozen=True) +class _TensorIdentity: + object_id: int + pointer: int + shape: tuple[int, ...] + stride: tuple[int, ...] + dtype: str + element_size: int + device: int + name: str + + @classmethod + def capture(cls, value: object, *, name: str) -> "_TensorIdentity": + shape = _shape(value, name=name) + stride = _stride(value, name=name) + if len(shape) != len(stride): + raise ValueError(f"{name} shape/stride rank mismatch") + return cls( + object_id=id(value), + pointer=_pointer(value, name=name), + shape=shape, + stride=stride, + dtype=_dtype(value), + element_size=_element_size(value, name=name), + device=_device_index(value, name=name), + name=name, + ) + + @property + def metadata(self) -> tuple: + """The layout a view must match; deliberately excludes the pointer.""" + return (self.shape, self.stride, self.dtype, self.element_size, + self.device) + + +@dataclass(frozen=True) +class BoundLiveWeightArena: + """Validated pointer tables and strong references for one SAMI rank.""" + + arena: "LiveWeightArena" = field(repr=False, compare=False) + world: int + rank: int + home_count: int + helper_count: int + device: int + plane_names: tuple[str, ...] + plane_bytes: tuple[int, ...] + local_plane_views: tuple[object, ...] = field(repr=False, compare=False) + owner_source_ptrs: tuple[int, ...] + scatter_source_ptrs: tuple[int, ...] + destination_ptrs: tuple[int, ...] + terminal_local_ptr: int + terminal_mc_ptr: int + + @property + def local_slot_count(self) -> int: + return self.home_count + self.helper_count + + +@dataclass(frozen=True) +class LiveWeightArena: + """Seven plane-major live weight views plus one terminal allocation.""" + + planes: tuple[LivePlaneView, ...] + terminals: LiveTerminalView + + def bind( + self, + *, + world: int, + rank: int, + home_count: int, + helper_count: int, + bundle: BundleLayout, + device: int, + ) -> BoundLiveWeightArena: + world = _exact_positive_int("world", world) + home_count = _exact_positive_int("home_count", home_count) + helper_count = _exact_positive_int("helper_count", helper_count) + if type(rank) is not int or not 0 <= rank < world: + raise ValueError("rank must be in [0, world)") + if type(device) is not int or device < 0: + raise ValueError("device must be a non-negative exact int") + + expected_names = tuple(plane.name for plane in bundle.planes) + observed_names = tuple(plane.name for plane in self.planes) + if observed_names != expected_names: + raise ValueError( + f"live arena plane order mismatch: {observed_names} != {expected_names}" + ) + if len(self.planes) != 7: + raise ValueError("live arena must contain exactly seven planes") + + slot_count = home_count + helper_count + local_planes: list[object] = [] + plane_bytes: list[int] = [] + owner_sources: list[int] = [] + scatter_sources: list[int] = [] + destinations: list[int] = [] + uc_ranges: list[list[tuple[int, int, str]]] = [ + [] for _ in range(world) + ] + mc_ranges: list[tuple[int, int, str]] = [] + + plane_identities: list[tuple[_TensorIdentity, ...]] = [] + mc_identities: list[_TensorIdentity] = [] + for plane_index, (view, layout) in enumerate(zip(self.planes, bundle.planes)): + if view.aliases_same_backing is not True: + raise ValueError(f"{view.name} UC/MC alias is not provider-attested") + if len(view.uc_views) != world: + raise ValueError(f"{view.name} must expose exactly world UC views") + + captured = tuple( + _TensorIdentity.capture(value, name=f"{view.name}.uc[{peer}]") + for peer, value in enumerate(view.uc_views) + ) + mc = _TensorIdentity.capture(view.mc_view, name=f"{view.name}.mc") + reference = captured[rank] + ( + expected_shape, + expected_stride, + expected_dtype, + expected_element_size, + expected_pointer_alignment, + ) = layout.kernel_metadata(slot_count) + expected_metadata = ( + expected_shape, + expected_stride, + expected_dtype, + expected_element_size, + device, + ) + for identity in (*captured, mc): + if identity.metadata != expected_metadata: + raise ValueError( + f"{view.name} must have exact MegaMoE " + f"shape={expected_shape}, stride={expected_stride}, " + f"dtype={expected_dtype} on cuda:{device}" + ) + if identity.pointer % expected_pointer_alignment: + raise ValueError( + f"{view.name} pointer must be " + f"{expected_pointer_alignment}-byte aligned" + ) + if len({identity.pointer for identity in captured}) != world: + raise ValueError(f"{view.name} UC rank mappings must have distinct pointers") + if mc.pointer in {identity.pointer for identity in captured}: + raise ValueError(f"{view.name} MC and UC mappings must be distinct") + if layout.nbytes % reference.element_size: + raise ValueError(f"{view.name} bytes/expert are not dtype aligned") + + expert_elements = layout.nbytes // reference.element_size + tail_span = 1 + sum( + (extent - 1) * stride + for extent, stride in zip(reference.shape[1:], reference.stride[1:]) + ) + if reference.stride[0] != expert_elements or tail_span != expert_elements: + raise ValueError( + f"{view.name} must be dense per expert with byte stride {layout.nbytes}" + ) + + plane_identities.append(captured) + mc_identities.append(mc) + local_planes.append(view.uc_views[rank]) + plane_bytes.append(layout.nbytes) + for peer, identity in enumerate(captured): + uc_ranges[peer].append( + ( + identity.pointer, + identity.pointer + slot_count * layout.nbytes, + view.name, + ) + ) + mc_ranges.append( + (mc.pointer, mc.pointer + slot_count * layout.nbytes, view.name) + ) + + weight_ranges = [ + (begin, end, f"{name}.uc[{peer}]") + for peer, peer_ranges in enumerate(uc_ranges) + for begin, end, name in peer_ranges + ] + weight_ranges.extend( + (begin, end, f"{name}.mc") for begin, end, name in mc_ranges + ) + self._reject_overlaps(weight_ranges) + + for local_expert in range(home_count): + for plane_index, layout in enumerate(bundle.planes): + owner_sources.append( + plane_identities[plane_index][rank].pointer + + local_expert * layout.nbytes + ) + for expert in range(world * home_count): + owner, local_expert = divmod(expert, home_count) + for plane_index, layout in enumerate(bundle.planes): + scatter_sources.append( + plane_identities[plane_index][owner].pointer + + local_expert * layout.nbytes + ) + for helper_index in range(helper_count): + physical_slot = home_count + helper_index + for plane_index, layout in enumerate(bundle.planes): + destinations.append( + mc_identities[plane_index].pointer + + physical_slot * layout.nbytes + ) + + if self.terminals.aliases_same_backing is not True: + raise ValueError("terminal UC/MC alias is not provider-attested") + terminal_local = _TensorIdentity.capture( + self.terminals.local_view, name="terminal.local" + ) + terminal_mc = _TensorIdentity.capture(self.terminals.mc_view, name="terminal.mc") + for terminal in (terminal_local, terminal_mc): + if terminal.shape != (world,) or terminal.stride != (1,): + raise ValueError("terminal views must have contiguous shape [EP]") + if terminal.dtype != "uint64" or terminal.element_size != 8: + raise TypeError("terminal views must have dtype uint64") + if terminal.device != device: + raise ValueError("terminal views are on the wrong CUDA device") + if terminal.pointer % 8: + raise ValueError("terminal pointers must be 8-byte aligned") + self._reject_terminal_overlap(terminal_local, weight_ranges, "local") + self._reject_terminal_overlap(terminal_mc, weight_ranges, "multicast") + terminal_bytes = world * 8 + if ( + terminal_local.pointer < terminal_mc.pointer + terminal_bytes + and terminal_mc.pointer < terminal_local.pointer + terminal_bytes + ): + raise ValueError("terminal local and multicast mappings overlap") + + return BoundLiveWeightArena( + arena=self, + world=world, + rank=rank, + home_count=home_count, + helper_count=helper_count, + device=device, + plane_names=expected_names, + plane_bytes=tuple(plane_bytes), + local_plane_views=tuple(local_planes), + owner_source_ptrs=tuple(owner_sources), + scatter_source_ptrs=tuple(scatter_sources), + destination_ptrs=tuple(destinations), + terminal_local_ptr=terminal_local.pointer, + terminal_mc_ptr=terminal_mc.pointer, + ) + + @staticmethod + def _reject_overlaps(ranges: list[tuple[int, int, str]]) -> None: + ordered = sorted(ranges) + for (_, previous_end, previous), (begin, _, current) in zip( + ordered, ordered[1:] + ): + if begin < previous_end: + raise ValueError( + f"live arena mappings overlap: {previous} and {current}" + ) + + @staticmethod + def _reject_terminal_overlap( + terminal: _TensorIdentity, + ranges: list[tuple[int, int, str]], + label: str, + ) -> None: + begin, end = terminal.pointer, terminal.pointer + terminal.shape[0] * 8 + for plane_begin, plane_end, name in ranges: + if begin < plane_end and plane_begin < end: + raise ValueError( + f"{label} terminal overlaps live weight plane {name}" + ) + + +@runtime_checkable +class LiveWeightArenaProvider(Protocol): + """Provider boundary implemented by TEKit or an equivalent allocator.""" + + def build_live_arena( + self, + *, + world: int, + rank: int, + home_count: int, + helper_count: int, + bundle: BundleLayout, + device: int, + comm: object, + ) -> LiveWeightArena: + """Return one fully constructed and provider-attested live arena.""" + + +__all__ = [ + "BoundLiveWeightArena", + "LivePlaneView", + "LiveTerminalView", + "LiveWeightArena", + "LiveWeightArenaProvider", +] diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/fabric.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/fabric.py new file mode 100644 index 000000000000..a785bc2690ba --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/fabric.py @@ -0,0 +1,900 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CUDA Fabric and NVSwitch multicast memory. + +This module is in the runtime dependency closure. It is imported lazily by the +live-weight arena to avoid eager native initialization at module import time. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from typing import Callable + +from cuda.bindings import driver as cuda + +from ._util import check_cuda + + +FABRIC = cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC + + +def _align_up(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +def _read_write_access(device: int) -> cuda.CUmemAccessDesc: + descriptor = cuda.CUmemAccessDesc() + descriptor.location.type = cuda.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + descriptor.location.id = int(device) + descriptor.flags = cuda.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE + return descriptor + + +def _allocation_properties(device: int) -> cuda.CUmemAllocationProp: + properties = cuda.CUmemAllocationProp() + properties.type = cuda.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED + properties.location.type = cuda.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + properties.location.id = int(device) + properties.requestedHandleTypes = FABRIC + return properties + + +@dataclass +class FabricAllocation: + """One exported device allocation plus the peer mappings imported from it.""" + + device: int + size: int + granularity: int + handle: object + local_ptr: int + shareable: bytes + _imported_handles: list[object] = field(default_factory=list, repr=False) + + @classmethod + def create( + cls, nbytes: int, device: int, *, extra_alignment: int = 1 + ) -> "FabricAllocation": + if nbytes <= 0: + raise ValueError("allocation size must be positive") + properties = _allocation_properties(device) + granularity = int( + check_cuda( + cuda.cuMemGetAllocationGranularity( + properties, + cuda.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED, + ), + "cuMemGetAllocationGranularity", + ) + ) + alignment = math.lcm(granularity, max(1, int(extra_alignment))) + size = _align_up(nbytes, alignment) + handle = check_cuda(cuda.cuMemCreate(size, properties, 0), "cuMemCreate") + allocation = cls(device, size, alignment, handle, 0, b"") + allocation.local_ptr = allocation._map(handle) + check_cuda( + cuda.cuMemsetD8(cuda.CUdeviceptr(allocation.local_ptr), 0, size), + "cuMemsetD8", + ) + exported = check_cuda( + cuda.cuMemExportToShareableHandle(handle, FABRIC, 0), + "cuMemExportToShareableHandle", + ) + allocation.shareable = bytes(exported.data) + if len(allocation.shareable) != 64: + raise RuntimeError( + f"fabric handle has {len(allocation.shareable)} bytes; expected 64" + ) + return allocation + + def _map(self, handle: object) -> int: + address = check_cuda( + cuda.cuMemAddressReserve(self.size, self.granularity, 0, 0), + "cuMemAddressReserve", + ) + check_cuda(cuda.cuMemMap(address, self.size, 0, handle, 0), "cuMemMap") + check_cuda( + cuda.cuMemSetAccess( + address, + self.size, + [_read_write_access(self.device)], + 1, + ), + "cuMemSetAccess", + ) + return int(address) + + def import_peer(self, shareable: bytes) -> int: + if len(shareable) != 64: + raise ValueError("peer fabric handle must contain exactly 64 bytes") + handle = check_cuda( + cuda.cuMemImportFromShareableHandle(bytearray(shareable), FABRIC), + "cuMemImportFromShareableHandle", + ) + self._imported_handles.append(handle) + return self._map(handle) + + +@dataclass(frozen=True) +class MulticastGroup: + """All-member NVSwitch multicast aperture bound to each rank's backing.""" + + device: int + size: int + granularity: int + handle: object + mc_ptr: int + creator_rank: int + backing_owner: FabricAllocation = field(repr=False, compare=False) + + alias_set: "MulticastTeamPair | None" = field(default=None, repr=False, compare=False) + + @property + def aliases(self) -> tuple: + return self.alias_set.aliases if self.alias_set is not None else () + + @property + def closed(self) -> bool: + return self.alias_set.closed if self.alias_set is not None else False + + def metadata(self) -> dict: + if self.alias_set is not None: + return self.alias_set.metadata() + return {"device": self.device, "size": self.size, + "creator_rank": self.creator_rank, "mc_ptr": self.mc_ptr, + "alias_count": 1, "dual_creator": False, "closed": False} + + def close_collectively(self) -> bool: + """Close a managed default group on every EP rank via its owners list. + + Explicit-creator legacy groups retain their original lifetime contract. + """ + if self.alias_set is None: + raise RuntimeError("explicit-creator groups have no managed pair lifecycle") + success = self.alias_set.close_collectively() + if success: + object.__setattr__(self, "mc_ptr", 0) + object.__setattr__(self, "handle", None) + return success + + @staticmethod + def granularity_for(world: int, nbytes: int, device: int) -> tuple[int, int]: + properties = cuda.CUmulticastObjectProp() + properties.numDevices = world + properties.size = nbytes + properties.handleTypes = int(FABRIC) + properties.flags = 0 + multicast = int( + check_cuda( + cuda.cuMulticastGetGranularity( + properties, + cuda.CUmulticastGranularity_flags.CU_MULTICAST_GRANULARITY_RECOMMENDED, + ), + "cuMulticastGetGranularity", + ) + ) + allocation = int( + check_cuda( + cuda.cuMemGetAllocationGranularity( + _allocation_properties(device), + cuda.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED, + ), + "cuMemGetAllocationGranularity", + ) + ) + return multicast, allocation + + @classmethod + def create( + cls, + *, + world: int, + rank: int, + creator_rank: int, + device: int, + size: int, + backing: FabricAllocation, + broadcast: Callable[[object, int], object], + barrier: Callable[[], None], + ) -> "MulticastGroup": + properties = cuda.CUmulticastObjectProp() + properties.numDevices = world + properties.size = size + properties.handleTypes = int(FABRIC) + properties.flags = 0 + granularity = int( + check_cuda( + cuda.cuMulticastGetGranularity( + properties, + cuda.CUmulticastGranularity_flags.CU_MULTICAST_GRANULARITY_RECOMMENDED, + ), + "cuMulticastGetGranularity", + ) + ) + if size % granularity: + raise ValueError("multicast size is not aligned to its granularity") + if backing.size < size: + raise ValueError("backing allocation is smaller than multicast object") + + serialized = None + if rank == creator_rank: + handle = check_cuda( + cuda.cuMulticastCreate(properties), "cuMulticastCreate" + ) + exported = check_cuda( + cuda.cuMemExportToShareableHandle(handle, FABRIC, 0), + "cuMemExportToShareableHandle(multicast)", + ) + serialized = bytes(exported.data) + serialized = broadcast(serialized, creator_rank) + if rank != creator_rank: + handle = check_cuda( + cuda.cuMemImportFromShareableHandle(bytearray(serialized), FABRIC), + "cuMemImportFromShareableHandle(multicast)", + ) + + check_cuda(cuda.cuMulticastAddDevice(handle, device), "cuMulticastAddDevice") + barrier() + check_cuda( + cuda.cuMulticastBindMem_v2( + handle, device, 0, backing.handle, 0, size, 0 + ), + "cuMulticastBindMem_v2", + ) + address = check_cuda( + cuda.cuMemAddressReserve(size, granularity, 0, 0), + "cuMemAddressReserve(multicast)", + ) + check_cuda(cuda.cuMemMap(address, size, 0, handle, 0), "cuMemMap(multicast)") + check_cuda( + cuda.cuMemSetAccess( + address, + size, + [_read_write_access(device)], + 1, + ), + "cuMemSetAccess(multicast)", + ) + barrier() + return cls( + device, + size, + granularity, + handle, + int(address), + creator_rank, + backing, + ) + + @classmethod + def create_for_members( + cls, *, world: int, members: tuple[int, ...], rank: int, + creator_rank: int | None = None, device: int, size: int, + backing: FabricAllocation, broadcast: Callable[[object, int], object], + barrier: Callable[[], None], all_gather: Callable | None = None, + owners: list | None = None, + ) -> "MulticastGroup | None": + """Default to two creators sharing backing; preserve explicit placement. + + The first member uses the second member's alias. Every other member + uses the first member's alias. A singleton creates one object. Pass an + ``owners`` list on every EP rank to retain even nonmember lifecycle + tokens; close these pairs in reverse global creation order. + """ + if creator_rank is not None: + return cls._create_single_for_members(world=world, members=members, + rank=rank, creator_rank=creator_rank, device=device, size=size, + backing=backing, broadcast=broadcast, barrier=barrier) + # These two values define the fallback control communicator itself. + if type(world) is not int or world <= 0: + raise ValueError("world must be a positive exact int") + if type(rank) is not int or not 0 <= rank < world: + raise ValueError("rank must be in [0, world)") + if (type(members) is not tuple or not members or + any(type(member) is not int for member in members)): + raise ValueError("multicast members must be a non-empty tuple of ranks") + count, begin = len(members), members[0] + if members != tuple(range(begin, begin + count)): + raise ValueError("multicast members must be contiguous sorted ranks") + if begin < 0 or members[-1] >= world: + raise ValueError("multicast members must be in [0, world)") + if world % count or begin % count: + raise ValueError("multicast members must form an aligned EP subgroup") + if type(device) is not int or device < 0: + raise ValueError("device must be a non-negative exact int") + if type(size) is not int or size <= 0: + raise ValueError("multicast size must be a positive exact int") + if owners is not None and type(owners) is not list: + raise TypeError("owners must be a list") + if not callable(broadcast) or not callable(barrier): + raise TypeError("broadcast and barrier must be callable") + if all_gather is None: + def all_gather(value): + return [broadcast(value if rank == root else None, root) + for root in range(world)] + if not callable(all_gather): + raise TypeError("all_gather must be callable") + + def validate_aligned_group(): + if not isinstance(backing, FabricAllocation): + raise TypeError("backing must be a FabricAllocation") + if backing.device != device or backing.size < size: + raise ValueError("backing device or size differs from multicast") + if rank in members: + granularity = cls.granularity_for(count, size, device)[0] + if granularity <= 0 or size % granularity: + raise ValueError("subgroup multicast size is not granular") + + try: + pair = MulticastTeam.create_dual(world=world, rank=rank, device=device, + members=members, size=size, all_gather=all_gather, + backing=backing if rank in members else None, + _validate=validate_aligned_group) + except MulticastTeamError as exc: + if owners is not None: + owners.append(exc.team) + raise + if owners is not None: + owners.append(pair) + if rank not in members: + return None + return cls(device, size, pair.granularity, pair.handle, pair.mc_ptr, + pair.creator_rank, backing, pair) + + @classmethod + def _create_single_for_members( + cls, + *, + world: int, + members: tuple[int, ...], + rank: int, + creator_rank: int | None = None, + device: int, + size: int, + backing: FabricAllocation, + broadcast: Callable[[object, int], object], + barrier: Callable[[], None], + ) -> "MulticastGroup | None": + """Collectively create one aligned subgroup multicast alias. + + Every EP rank calls this method in the same order so setup may keep its + existing world control communicator. Only ``members`` add a device, + bind their local backing, and map the multicast virtual address. + An omitted creator selects the second member (the sole member for a + singleton); an explicit creator preserves caller placement. This is + a deterministic placement policy, not a topology-independent guarantee. + """ + + if type(world) is not int or world <= 0: + raise ValueError("world must be a positive exact int") + if type(rank) is not int or not 0 <= rank < world: + raise ValueError("rank must be in [0, world)") + if ( + not isinstance(members, tuple) + or not members + or any(type(member) is not int for member in members) + ): + raise ValueError("multicast members must be a non-empty tuple of ranks") + group_size = len(members) + group_begin = members[0] + if members != tuple(range(group_begin, group_begin + group_size)): + raise ValueError("multicast members must be contiguous sorted ranks") + if group_begin < 0 or members[-1] >= world: + raise ValueError("multicast members must be in [0, world)") + if world % group_size or group_begin % group_size: + raise ValueError("multicast members must form an aligned EP subgroup") + if creator_rank is None: + creator_rank = members[1] if group_size > 1 else members[0] + if type(creator_rank) is not int: + raise ValueError("multicast creator must be an exact rank") + if creator_rank not in members: + raise ValueError("multicast creator must belong to the subgroup") + if type(device) is not int or device < 0: + raise ValueError("device must be a non-negative exact int") + if type(size) is not int or size <= 0: + raise ValueError("multicast size must be a positive exact int") + if not isinstance(backing, FabricAllocation): + raise TypeError("backing must be a FabricAllocation") + if backing.device != device: + raise ValueError("backing and multicast devices differ") + if not callable(broadcast) or not callable(barrier): + raise TypeError("broadcast and barrier must be callable") + properties = cuda.CUmulticastObjectProp() + properties.numDevices = len(members) + properties.size = size + properties.handleTypes = int(FABRIC) + properties.flags = 0 + granularity = int( + check_cuda( + cuda.cuMulticastGetGranularity( + properties, + cuda.CUmulticastGranularity_flags.CU_MULTICAST_GRANULARITY_RECOMMENDED, + ), + "cuMulticastGetGranularity(subgroup)", + ) + ) + if size % granularity: + raise ValueError("subgroup multicast size is not granular") + if backing.size < size: + raise ValueError("backing allocation is smaller than subgroup multicast") + + serialized = None + handle = None + if rank == creator_rank: + handle = check_cuda( + cuda.cuMulticastCreate(properties), + "cuMulticastCreate(subgroup)", + ) + exported = check_cuda( + cuda.cuMemExportToShareableHandle(handle, FABRIC, 0), + "cuMemExportToShareableHandle(subgroup)", + ) + serialized = bytes(exported.data) + serialized = broadcast(serialized, creator_rank) + try: + serialized = bytes(serialized) + except (TypeError, ValueError) as exc: + raise TypeError("broadcast returned an invalid multicast handle") from exc + if len(serialized) != 64: + raise RuntimeError( + f"subgroup fabric handle has {len(serialized)} bytes; expected 64" + ) + is_member = rank in members + if is_member and rank != creator_rank: + handle = check_cuda( + cuda.cuMemImportFromShareableHandle(bytearray(serialized), FABRIC), + "cuMemImportFromShareableHandle(subgroup)", + ) + if is_member: + check_cuda( + cuda.cuMulticastAddDevice(handle, device), + "cuMulticastAddDevice(subgroup)", + ) + barrier() + address = 0 + if is_member: + check_cuda( + cuda.cuMulticastBindMem_v2( + handle, device, 0, backing.handle, 0, size, 0 + ), + "cuMulticastBindMem_v2(subgroup)", + ) + address = check_cuda( + cuda.cuMemAddressReserve(size, granularity, 0, 0), + "cuMemAddressReserve(subgroup)", + ) + check_cuda( + cuda.cuMemMap(address, size, 0, handle, 0), + "cuMemMap(subgroup)", + ) + check_cuda( + cuda.cuMemSetAccess(address, size, [_read_write_access(device)], 1), + "cuMemSetAccess(subgroup)", + ) + barrier() + if not is_member: + return None + return cls( + device, + size, + granularity, + handle, + int(address), + creator_rank, + backing, + ) + + +def pick_multicast_creator(rank_to_node: list[int], source_rank: int = 0) -> int: + """Select a creator off source rank's node when the group spans nodes.""" + source_node = rank_to_node[source_rank] + for rank, node in enumerate(rank_to_node): + if node != source_node: + return rank + return source_rank + + +__all__ = [ + "FabricAllocation", + "MulticastGroup", + "check_cuda", + "pick_multicast_creator", +] + + +_RETAINED_MULTICAST_TEAMS: list[object] = [] + + +class MulticastTeamError(RuntimeError): + """A collective team operation failed; keep its partial resource graph.""" + + def __init__(self, message: str, team: "MulticastTeam", failures: list) -> None: + super().__init__(message) + self.team, self.failures = team, failures + + +class MulticastTeam: + """Explicit arbitrary-member multicast team with collective teardown. + + Every EP rank participates in identical create/close calls. Members with a + backing borrow it; members without one own a fresh dummy. Nonmembers never + allocate or map. Callers must release consumers before collective close. + There is deliberately no destructor that starts a distributed collective. + """ + + def __init__(self, *, world, rank, device, members, creator_rank, size, + all_gather, backing) -> None: + self.world, self.rank, self.device = world, rank, device + self.members, self.creator_rank, self.size = members, creator_rank, size + self.all_gather, self.backing_owner = all_gather, backing + self._borrowed = backing is not None + self.is_member = self.bound = self.mc_mapped = self._uc_mapped = False + self.handle = self._owned_handle = None + self.mc_ptr = self.local_ptr = self._owned_ptr = 0 + self.closed = self.retained_after_cleanup_failure = False + self.granularities = None + + @property + def owns_backing(self) -> bool: + return self.is_member and not self._borrowed + + def _call(self, name, *args): + return check_cuda(getattr(cuda, name)(*args), name) + + def _phase(self, label, operation, *, cleanup=False): + value, error = None, None + try: + value = operation() + except Exception as exc: + error = {"rank": self.rank, "type": type(exc).__name__, "message": str(exc)} + rows = self.all_gather({"rank": self.rank, "world": self.world, + "phase": label, "error": error}) + failures = [row["error"] for row in rows if row["error"] is not None] + if ([row["rank"] for row in rows] != list(range(len(rows))) or + any(row["world"] != len(rows) or row["phase"] != label for row in rows)): + failures.append({"message": "multicast team collective phase/rank mismatch"}) + if failures and not cleanup: + raise MulticastTeamError(label + ": " + repr(failures), self, failures) + return not failures if cleanup else value + + @staticmethod + def granularities_for(member_count: int, nbytes: int, device: int) -> dict: + """Query before allocation; sizes are never silently enlarged by create.""" + if any(type(x) is not int for x in (member_count, nbytes, device)) or min(member_count, nbytes) <= 0 or device < 0: + raise ValueError("team geometry requires positive exact counts and a nonnegative device") + supported = check_cuda(cuda.cuDeviceGetAttribute( + cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED, device), + "cuDeviceGetAttribute(multicast)") + if int(supported) != 1: + raise RuntimeError("team device does not support multicast") + prop = cuda.CUmulticastObjectProp() + prop.numDevices, prop.size, prop.handleTypes, prop.flags = member_count, nbytes, int(FABRIC), 0 + cm, ca = cuda.CUmulticastGranularity_flags, cuda.CUmemAllocationGranularity_flags + values = {"mc_minimum": int(check_cuda(cuda.cuMulticastGetGranularity(prop, cm.CU_MULTICAST_GRANULARITY_MINIMUM))), + "mc_recommended": int(check_cuda(cuda.cuMulticastGetGranularity(prop, cm.CU_MULTICAST_GRANULARITY_RECOMMENDED))), + "allocation_minimum": int(check_cuda(cuda.cuMemGetAllocationGranularity(_allocation_properties(device), ca.CU_MEM_ALLOC_GRANULARITY_MINIMUM))), + "allocation_recommended": int(check_cuda(cuda.cuMemGetAllocationGranularity(_allocation_properties(device), ca.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED)))} + if any(x <= 0 or x & (x - 1) for x in values.values()): + raise ValueError("invalid CUDA team granularity") + return values + + def _validate(self): + if type(self.world) is not int or self.world <= 0: + raise ValueError("world must be a positive exact int") + if type(self.rank) is not int or not 0 <= self.rank < self.world: + raise ValueError("rank must be an exact EP rank") + if type(self.device) is not int or self.device < 0: + raise ValueError("device must be a nonnegative exact ordinal") + if (type(self.members) is not tuple or not self.members or + any(type(x) is not int for x in self.members) or + self.members != tuple(sorted(set(self.members))) or + any(not 0 <= x < self.world for x in self.members)): + raise ValueError("members must be sorted unique exact ranks within EP") + if type(self.creator_rank) is not int or self.creator_rank not in self.members: + raise ValueError("creator must be an exact member rank") + if type(self.size) is not int or self.size <= 0: + raise ValueError("size must be a positive exact byte count") + self.is_member = self.rank in self.members + if not self.is_member and self.backing_owner is not None: + raise ValueError("nonmembers cannot provide backing") + if self.backing_owner is not None: + if not isinstance(self.backing_owner, FabricAllocation): + raise TypeError("borrowed backing must be a FabricAllocation") + if self.backing_owner.device != self.device or self.backing_owner.size < self.size or not self.backing_owner.local_ptr: + raise ValueError("borrowed backing device, size or UC mapping mismatch") + self.local_ptr = int(self.backing_owner.local_ptr) + + def _initialize(self): + self._phase("validate team arguments", self._validate) + configs = self.all_gather((self.world, self.members, self.creator_rank, self.size)) + def uniform(): + if any(row != configs[0] for row in configs): + raise ValueError("multicast team configuration differs across ranks") + self._phase("uniform team configuration", uniform) + def granularity(): + if self.is_member: + self.granularities = self.granularities_for(len(self.members), self.size, self.device) + if self.size % self.granularities["mc_minimum"] or self.size % self.granularities["allocation_minimum"]: + raise ValueError("explicit size is not granular; no automatic enlargement") + return self.granularities + values = self.all_gather(self._phase("team capability and granularity", granularity)) + self.granularity = max(max(row.values()) for row in values if row is not None) + def own_backing(): + if self.owns_backing: + self._owned_handle = self._call("cuMemCreate", self.size, _allocation_properties(self.device), 0) + self.backing_owner = FabricAllocation(self.device, self.size, self.granularity, self._owned_handle, 0, b"") + self._phase("fresh or borrowed backing", own_backing) + def reserve_uc(): + if self.owns_backing: + self._owned_ptr = int(self._call("cuMemAddressReserve", self.size, self.granularity, 0, 0)) + self._phase("owned UC reservation", reserve_uc) + def map_uc(): + if self.owns_backing: + self._call("cuMemMap", self._owned_ptr, self.size, 0, self._owned_handle, 0) + self._uc_mapped = True + self._phase("owned UC mapping", map_uc) + def access_uc(): + if self.owns_backing: + self._call("cuMemSetAccess", self._owned_ptr, self.size, [_read_write_access(self.device)], 1) + self.local_ptr = self._owned_ptr + self.backing_owner.local_ptr = self.local_ptr + self._phase("owned UC access", access_uc) + def create(): + if self.rank == self.creator_rank: + prop = cuda.CUmulticastObjectProp() + prop.numDevices, prop.size, prop.handleTypes, prop.flags = len(self.members), self.size, int(FABRIC), 0 + self.handle = self._call("cuMulticastCreate", prop) + self._phase("creator MC creation", create) + def export(): + if self.rank == self.creator_rank: + data = bytes(self._call("cuMemExportToShareableHandle", self.handle, FABRIC, 0).data) + if len(data) != 64: + raise ValueError("FABRIC handle must contain exactly 64 bytes") + return data + exports = self.all_gather(self._phase("creator MC export", export)) + def import_member(): + data = exports[self.creator_rank] + if type(data) is not bytes or len(data) != 64: + raise ValueError("invalid collective FABRIC handle") + if self.is_member and self.rank != self.creator_rank: + self.handle = self._call("cuMemImportFromShareableHandle", bytearray(data), FABRIC) + self._phase("member MC import", import_member) + self._phase("all AddDevice before bind/map", lambda: + self._call("cuMulticastAddDevice", self.handle, self.device) if self.is_member else None) + def bind(): + if self.is_member: + self._call("cuMulticastBindMem_v2", self.handle, self.device, 0, self.backing_owner.handle, 0, self.size, 0) + self.bound = True + self._phase("member backing bind", bind) + def reserve_mc(): + if self.is_member: + self.mc_ptr = int(self._call("cuMemAddressReserve", self.size, self.granularity, 0, 0)) + self._phase("member MC reservation", reserve_mc) + def map_mc(): + if self.is_member: + self._call("cuMemMap", self.mc_ptr, self.size, 0, self.handle, 0) + self.mc_mapped = True + self._phase("member MC mapping", map_mc) + self._phase("member MC access", lambda: + self._call("cuMemSetAccess", self.mc_ptr, self.size, [_read_write_access(self.device)], 1) if self.is_member else None) + return self + + @classmethod + def create(cls, *, world, rank, device, members, creator_rank, size, + all_gather, backing=None) -> "MulticastTeam": + if not callable(all_gather): + raise TypeError("all_gather must be callable") + team = cls(world=world, rank=rank, device=device, members=members, + creator_rank=creator_rank, size=size, all_gather=all_gather, backing=backing) + try: + return team._initialize() + except Exception as exc: + try: + team.close_collectively() + except Exception: + team._retain() + if isinstance(exc, MulticastTeamError): + raise + raise MulticastTeamError(str(exc), team, [{"message": str(exc)}]) from exc + + @classmethod + def create_dual(cls, *, world, rank, device, members, size, all_gather, + backing=None, _validate=None) -> "MulticastTeamPair": + """Create first/second-member aliases, with at most one owned backing. + + Explicit ``create(..., creator_rank=...)`` semantics are unchanged. + """ + return MulticastTeamPair.create(world=world, rank=rank, device=device, + members=members, size=size, all_gather=all_gather, backing=backing, + team_factory=cls, validate=_validate) + + def _retain(self): + if not self.retained_after_cleanup_failure: + _RETAINED_MULTICAST_TEAMS.append(self) + self.retained_after_cleanup_failure = True + + def close_collectively(self) -> bool: + """Drain and free in dependency order; stop collectively on any error.""" + if self.closed: + return True + def unmap_mc(): + if self.mc_mapped: + self._call("cuMemUnmap", self.mc_ptr, self.size) + self.mc_mapped = False + if self.mc_ptr: + self._call("cuMemAddressFree", self.mc_ptr, self.size) + self.mc_ptr = 0 + def unbind(): + if self.bound: + self._call("cuMulticastUnbind", self.handle, self.device, 0, self.size) + self.bound = False + def release_imports(): + if self.rank != self.creator_rank and self.handle is not None: + self._call("cuMemRelease", self.handle) + self.handle = None + def release_creator(): + if self.rank == self.creator_rank and self.handle is not None: + self._call("cuMemRelease", self.handle) + self.handle = None + def unmap_uc(): + if self._uc_mapped: + self._call("cuMemUnmap", self._owned_ptr, self.size) + self._uc_mapped = False + if self._owned_ptr: + self._call("cuMemAddressFree", self._owned_ptr, self.size) + self._owned_ptr = self.local_ptr = 0 + def release_backing(): + if self._owned_handle is not None: + self._call("cuMemRelease", self._owned_handle) + self._owned_handle = None + self.backing_owner = None + try: + for label, fn in (("drain context", lambda: self._call("cuCtxSynchronize")), + ("unmap MC", unmap_mc), ("unbind MC", unbind), + ("release imports", release_imports), ("release creator", release_creator), + ("unmap owned UC", unmap_uc), ("release owned backing", release_backing)): + if not self._phase("cleanup " + label, fn, cleanup=True): + self._retain() + return False + except Exception: + self._retain() + raise + self.closed = True + return True + + def metadata(self) -> dict: + return {"rank": self.rank, "device": self.device, "members": list(self.members), + "creator_rank": self.creator_rank, "size": self.size, "is_member": self.is_member, + "owns_backing": self.owns_backing, "borrowed_backing": self._borrowed, + "mc_ptr": self.mc_ptr, "local_ptr": self.local_ptr, "bound": self.bound, + "granularities": self.granularities, "closed": self.closed, + "retained_after_cleanup_failure": self.retained_after_cleanup_failure} + + +class MulticastTeamPair: + """Two MC objects over one backing, exposing the rank-selected aperture. + + The first alias owns any dummy; the second borrows that exact allocation. + Every EP rank retains this object, including nonmembers, and calls close in + the same global order. No destructor starts a collective. + """ + + _SELECTED_FIELDS = frozenset(("handle", "mc_ptr", "creator_rank", "bound", + "mc_mapped", "granularity", "granularities")) + _PRIMARY_FIELDS = frozenset(("backing_owner", "owns_backing", "local_ptr", + "_owned_handle", "_owned_ptr", "_uc_mapped", "_borrowed", "is_member")) + + def __init__(self, *, world, rank, device, members, size, all_gather): + self.world, self.rank, self.device = world, rank, device + self.members, self.size, self.all_gather = members, size, all_gather + self._aliases = [] + self.closed = self.retained_after_cleanup_failure = False + + @property + def aliases(self) -> tuple[MulticastTeam, ...]: + return tuple(self._aliases) + + @property + def selected_alias(self) -> MulticastTeam: + if not self._aliases: + raise RuntimeError("multicast pair has no initialized alias") + index = 1 if len(self._aliases) > 1 and self.rank == self.members[0] else 0 + return self._aliases[index] + + def __getattr__(self, name): + if name in self._SELECTED_FIELDS: + return getattr(self.selected_alias, name) + if name in self._PRIMARY_FIELDS and self._aliases: + return getattr(self._aliases[0], name) + raise AttributeError(name) + + def _phase(self, label, operation, *, cleanup=False): + return MulticastTeam._phase(self, label, operation, cleanup=cleanup) + + def _retain(self): + MulticastTeam._retain(self) + + @classmethod + def create(cls, *, world, rank, device, members, size, all_gather, + backing=None, team_factory=MulticastTeam, validate=None): + if not callable(all_gather): + raise TypeError("all_gather must be callable") + pair = cls(world=world, rank=rank, device=device, members=members, + size=size, all_gather=all_gather) + try: + creator = members[0] if type(members) is tuple and members else None + probe = team_factory(world=world, rank=rank, device=device, + members=members, creator_rank=creator, size=size, + all_gather=all_gather, backing=backing) + def validate_arguments(): + probe._validate() + if validate is not None: + validate() + pair._phase("validate dual multicast arguments", validate_arguments) + primary = team_factory.create(world=world, rank=rank, device=device, + members=members, creator_rank=members[0], size=size, + all_gather=all_gather, backing=backing) + pair._aliases.append(primary) + if len(members) > 1: + secondary = team_factory.create(world=world, rank=rank, device=device, + members=members, creator_rank=members[1], size=size, + all_gather=all_gather, + backing=primary.backing_owner if primary.is_member else None) + pair._aliases.append(secondary) + return pair + except Exception as exc: + if isinstance(exc, MulticastTeamError) and exc.team is not pair: + if all(alias is not exc.team for alias in pair._aliases): + pair._aliases.append(exc.team) + # A failed secondary cleanup may still reference primary's dummy. + # Never free the primary allocation while any such alias survives. + if any(alias.retained_after_cleanup_failure for alias in pair._aliases): + pair._retain() + else: + try: + if not pair.close_collectively(): + pair._retain() + except Exception: + pair._retain() + failures = exc.failures if isinstance(exc, MulticastTeamError) else [{"message": str(exc)}] + raise MulticastTeamError(str(exc), pair, failures) from exc + + def close_collectively(self) -> bool: + """Free the borrowed secondary first; the primary owns the only dummy.""" + if self.closed: + return True + if self.retained_after_cleanup_failure: + return False + try: + for alias in reversed(self._aliases): + if alias.retained_after_cleanup_failure or not alias.close_collectively(): + self._retain() + return False + except Exception: + self._retain() + raise + self.closed = True + return True + + def metadata(self) -> dict: + selected = self.selected_alias.metadata() if self._aliases else {} + selected.update({"rank": self.rank, "device": self.device, + "members": list(self.members) if type(self.members) is tuple else [], + "size": self.size, "dual_creator": len(self.members) > 1 if type(self.members) is tuple else False, + "selection_policy": "first member uses second creator; all others use first creator", + "alias_count": len(self._aliases), + "aliases": [alias.metadata() for alias in self._aliases], + "closed": self.closed, + "retained_after_cleanup_failure": self.retained_after_cleanup_failure, + "owns_backing": self._aliases[0].owns_backing if self._aliases else False, + "allocated_bytes": sum(alias.size for alias in self._aliases if alias._owned_handle is not None), + "mapping_count": sum(bool(alias.mc_ptr) for alias in self._aliases), + "reserved_va_bytes": sum((alias.size if alias.mc_ptr else 0) + + (alias.size if alias._owned_ptr else 0) for alias in self._aliases)}) + return selected + + +__all__.extend(["MulticastTeam", "MulticastTeamPair", "MulticastTeamError"]) diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/geometry.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/geometry.py new file mode 100644 index 000000000000..2cdc1caaa014 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/geometry.py @@ -0,0 +1,304 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Frozen seven-plane MegaMoE weight geometry. + +Storage placement is owned by :mod:`megamoe_scheduler.sami.arena`; production +uses one external plane-major live bank without internal receive storage. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + + +CANONICAL_WEIGHT_PLANE_NAMES = ( + "mega_fc1_weight", + "mega_fc1_weight_sf", + "mega_fc2_weight", + "mega_fc2_weight_sf", + "fc31_alpha", + "fc2_alpha", + "fc1_norm_const", +) + + +def _align_up(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +@dataclass(frozen=True) +class PlaneViewLayout: + """One per-expert tensor view of a live weight plane. + + Shape and stride omit the leading physical-slot axis. The descriptor is + Torch-independent, so validation does not resolve optional FP4/FP8 dtypes. + """ + + shape: tuple[int, ...] + stride: tuple[int, ...] + dtype: str + element_size: int + pointer_alignment: int + + def __post_init__(self) -> None: + if len(self.shape) != len(self.stride): + raise ValueError("plane view shape and stride must have equal rank") + if any(type(value) is not int or value <= 0 for value in self.shape): + raise ValueError("plane view extents must be positive exact integers") + if any(type(value) is not int or value <= 0 for value in self.stride): + raise ValueError("plane view strides must be positive exact integers") + if type(self.dtype) is not str or not self.dtype: + raise ValueError("plane view dtype must be a non-empty string") + if type(self.element_size) is not int or self.element_size <= 0: + raise ValueError("plane view element_size must be a positive exact integer") + if type(self.pointer_alignment) is not int or self.pointer_alignment <= 0: + raise ValueError("plane view pointer_alignment must be a positive exact integer") + + @property + def elements_per_slot(self) -> int: + if not self.shape: + return 1 + return 1 + sum((extent - 1) * stride for extent, stride in zip(self.shape, self.stride)) + + @property + def nbytes(self) -> int: + return self.elements_per_slot * self.element_size + + def batched_metadata( + self, slot_count: int + ) -> tuple[tuple[int, ...], tuple[int, ...], str, int, int]: + if type(slot_count) is not int or slot_count <= 0: + raise ValueError("slot_count must be a positive exact integer") + return ( + (slot_count, *self.shape), + (self.elements_per_slot, *self.stride), + self.dtype, + self.element_size, + self.pointer_alignment, + ) + + +@dataclass(frozen=True) +class Plane: + """One canonical MegaMoE weight plane within an expert bundle. + + The first four fields preserve the original public constructor. Layouts + created by BundleLayout.create additionally carry the storage view used by + the loader and the typed view consumed by the kernel. + """ + + name: str + offset: int + nbytes: int + alignment: int + storage: Optional[PlaneViewLayout] = None + kernel: Optional[PlaneViewLayout] = None + + def __post_init__(self) -> None: + if (self.storage is None) != (self.kernel is None): + raise ValueError("plane storage and kernel layouts must be provided together") + if self.storage is not None: + if self.storage.nbytes != self.nbytes: + raise ValueError( + f"{self.name} storage view spans {self.storage.nbytes} bytes; " + f"expected {self.nbytes}" + ) + assert self.kernel is not None + if self.kernel.nbytes != self.nbytes: + raise ValueError( + f"{self.name} kernel view spans {self.kernel.nbytes} bytes; " + f"expected {self.nbytes}" + ) + + def storage_metadata( + self, slot_count: int + ) -> tuple[tuple[int, ...], tuple[int, ...], str, int, int]: + if self.storage is None: + raise ValueError(f"{self.name} has no canonical storage view") + return self.storage.batched_metadata(slot_count) + + def kernel_metadata( + self, slot_count: int + ) -> tuple[tuple[int, ...], tuple[int, ...], str, int, int]: + if self.kernel is None: + raise ValueError(f"{self.name} has no canonical kernel view") + return self.kernel.batched_metadata(slot_count) + + +@dataclass(frozen=True) +class BundleLayout: + """NVFP4 bytes/expert for the seven plane-major live tensors. + + ``offset`` and ``stride_bytes`` describe a canonical packed serialization + only; direct-live SAMI uses each plane's ``nbytes`` as its physical-slot + byte stride and never allocates that serialization. + """ + + hidden: int + intermediate: int + planes: tuple[Plane, ...] + total_bytes: int + stride_bytes: int + alignment: int + expand_intermediate: Optional[int] = None + + @classmethod + def create( + cls, + hidden: int = 7168, + intermediate: int = 3072, + *, + expand_intermediate: Optional[int] = None, + weight_bits: int = 4, + sf_vec_size: int = 16, + sf_bits: int = 8, + plane_alignment: int = 512, + ) -> "BundleLayout": + dimensions = { + "hidden": hidden, + "intermediate": intermediate, + "weight_bits": weight_bits, + "sf_vec_size": sf_vec_size, + "sf_bits": sf_bits, + "plane_alignment": plane_alignment, + } + if any(type(value) is not int or value <= 0 for value in dimensions.values()): + raise ValueError("all bundle dimensions must be positive exact integers") + if expand_intermediate is None: + expand_intermediate = 2 * intermediate + if type(expand_intermediate) is not int or expand_intermediate <= 0: + raise ValueError("all bundle dimensions must be positive exact integers") + if (weight_bits, sf_vec_size, sf_bits) != (4, 16, 8): + raise ValueError( + "production bundle format is fixed to NVFP4 weights and " + "FP8 block-16 scale factors" + ) + if hidden % sf_vec_size or intermediate % sf_vec_size: + raise ValueError("hidden/intermediate must be block-scale aligned") + if expand_intermediate != 2 * intermediate: + raise ValueError( + "production bundle requires expand_intermediate == 2 * intermediate" + ) + + # These are the exact TEKit transformed-tensor formulas. The simpler + # elements/sf_vec_size expression only coincides at some model shapes. + fc1_sf_elements = _align_up(expand_intermediate, 128) * _align_up( + hidden // sf_vec_size, 4 + ) + fc2_sf_elements = _align_up(hidden, 128) * _align_up( + intermediate // sf_vec_size, 4 + ) + + def view( + shape: tuple[int, ...], + dtype: str, + element_size: int, + pointer_alignment: int, + *, + stride: Optional[tuple[int, ...]] = None, + ) -> PlaneViewLayout: + if stride is None: + steps = [] + running = 1 + for extent in reversed(shape): + steps.append(running) + running *= extent + stride = tuple(reversed(steps)) + return PlaneViewLayout( + shape=shape, + stride=stride, + dtype=dtype, + element_size=element_size, + pointer_alignment=pointer_alignment, + ) + + gate_up = expand_intermediate + specs = ( + ( + "mega_fc1_weight", + plane_alignment, + view((gate_up, hidden // 2), "uint8", 1, 16), + view( + (hidden // 2, gate_up), + "float4_e2m1fn_x2", + 1, + 16, + stride=(1, hidden // 2), + ), + ), + ( + "mega_fc1_weight_sf", + plane_alignment, + view((fc1_sf_elements,), "uint8", 1, 16), + view((fc1_sf_elements,), "float8_e4m3fn", 1, 16), + ), + ( + "mega_fc2_weight", + plane_alignment, + view((hidden, intermediate // 2), "uint8", 1, 16), + view( + (intermediate // 2, hidden), + "float4_e2m1fn_x2", + 1, + 16, + stride=(1, intermediate // 2), + ), + ), + ( + "mega_fc2_weight_sf", + plane_alignment, + view((fc2_sf_elements,), "uint8", 1, 16), + view((fc2_sf_elements,), "float8_e4m3fn", 1, 16), + ), + ( + "fc31_alpha", + 16, + view((), "float32", 4, 4), + view((), "float32", 4, 4), + ), + ( + "fc2_alpha", + 16, + view((), "float32", 4, 4), + view((), "float32", 4, 4), + ), + ( + "fc1_norm_const", + 16, + view((), "float32", 4, 4), + view((), "float32", 4, 4), + ), + ) + + planes: list[Plane] = [] + cursor = 0 + for name, alignment, storage, kernel in specs: + nbytes = storage.nbytes + if kernel.nbytes != nbytes: + raise ValueError(f"{name} storage and kernel views span different bytes") + cursor = _align_up(cursor, alignment) + planes.append(Plane(name, cursor, nbytes, alignment, storage=storage, kernel=kernel)) + cursor += nbytes + return cls( + hidden=hidden, + intermediate=intermediate, + planes=tuple(planes), + total_bytes=sum(plane.nbytes for plane in planes), + stride_bytes=_align_up(cursor, plane_alignment), + alignment=plane_alignment, + expand_intermediate=expand_intermediate, + ) + + @property + def plane_count(self) -> int: + return len(self.planes) + + +__all__ = [ + "BundleLayout", + "CANONICAL_WEIGHT_PLANE_NAMES", + "Plane", + "PlaneViewLayout", +] diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py new file mode 100644 index 000000000000..b6915a6fd3a7 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py @@ -0,0 +1,695 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Hierarchical HALO-Q live arena and low-overhead SAMI submission path.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import threading +from typing import Protocol, runtime_checkable + +import torch +from cuda.bindings import driver as cuda + +from .arena import BoundLiveWeightArena, LiveWeightArena, _TensorIdentity +from ._util import check_cuda +from ..geometry import hierarchy_group_sizes +from .geometry import BundleLayout + +@dataclass(frozen=True) +class BoundHierarchicalLiveWeightArena: + """Validated base arena plus this rank's multicast alias at every level.""" + + arena: "HierarchicalLiveWeightArena" = field(repr=False, compare=False) + base: BoundLiveWeightArena = field(repr=False, compare=False) + group_sizes: tuple[int, ...] + group_destination_ptrs: tuple[int, ...] + group_mc_views: tuple[tuple[object, ...], ...] = field( + repr=False, compare=False + ) + group_mc_owners: tuple[tuple[object, ...], ...] = field( + repr=False, compare=False + ) + + +@dataclass(frozen=True) +class HierarchicalLiveWeightArena: + """One external live bank with a multicast alias for each HALO level. + + ``group_mc_views`` and ``group_mc_owners`` are plane-major: + ``[plane][level]``. At a given level, each rank exposes only the alias and + strong VMM owner token for its own aligned group. Level zero is the + full-world alias already present in ``base.planes[*].mc_view``. + """ + + base: LiveWeightArena + group_sizes: tuple[int, ...] + group_mc_views: tuple[tuple[object, ...], ...] = field( + repr=False, compare=False + ) + group_mc_owners: tuple[tuple[object, ...], ...] = field( + repr=False, compare=False + ) + aliases_same_backing: bool = False + + def bind( + self, + *, + world: int, + rank: int, + home_count: int, + helper_count: int, + bundle: BundleLayout, + device: int, + ) -> BoundHierarchicalLiveWeightArena: + if self.aliases_same_backing is not True: + raise ValueError("hierarchical multicast aliases are not provider-attested") + expected_sizes = hierarchy_group_sizes(world) + if self.group_sizes != expected_sizes: + raise ValueError( + f"hierarchy group sizes {self.group_sizes} != {expected_sizes}" + ) + base = self.base.bind( + world=world, + rank=rank, + home_count=home_count, + helper_count=helper_count, + bundle=bundle, + device=device, + ) + if len(self.group_mc_views) != len(bundle.planes): + raise ValueError("hierarchical arena must expose every weight plane") + if len(self.group_mc_owners) != len(bundle.planes): + raise ValueError("hierarchical arena must retain every VMM owner") + if self.base.terminals.backing_owner is None: + raise ValueError("hierarchical terminal backing owner is not retained") + + identities: list[tuple[_TensorIdentity, ...]] = [] + for views, owners, base_view in zip( + self.group_mc_views, self.group_mc_owners, self.base.planes + ): + if base_view.backing_owner is None: + raise ValueError(f"{base_view.name} backing owner is not retained") + if len(views) != len(expected_sizes): + raise ValueError( + f"{base_view.name} has the wrong hierarchy depth" + ) + if len(owners) != len(expected_sizes) or any( + owner is None for owner in owners + ): + raise ValueError( + f"{base_view.name} must retain every subgroup VMM owner" + ) + captured = tuple( + _TensorIdentity.capture( + view, name=f"{base_view.name}.group_mc[{level}]" + ) + for level, view in enumerate(views) + ) + reference = _TensorIdentity.capture( + base_view.mc_view, name=f"{base_view.name}.mc" + ) + for identity in captured: + if identity.metadata != reference.metadata: + raise ValueError( + f"{base_view.name} group multicast metadata changed" + ) + if captured[0].pointer != reference.pointer: + raise ValueError( + f"{base_view.name} level-zero alias is not the base MC view" + ) + if len({identity.pointer for identity in captured}) != len(captured): + raise ValueError( + f"{base_view.name} hierarchy levels need distinct virtual mappings" + ) + identities.append(captured) + + destinations: list[int] = [] + for level in range(len(expected_sizes)): + for helper in range(helper_count): + physical_slot = home_count + helper + for plane_index, layout in enumerate(bundle.planes): + destinations.append( + identities[plane_index][level].pointer + + physical_slot * layout.nbytes + ) + return BoundHierarchicalLiveWeightArena( + arena=self, + base=base, + group_sizes=expected_sizes, + group_destination_ptrs=tuple(destinations), + group_mc_views=self.group_mc_views, + group_mc_owners=self.group_mc_owners, + ) + +# runtime_checkable because __init__ isinstance()-checks the provider below. +# Without it that check raises TypeError instead of validating anything, and +# only at runtime -- no import, compile or signature check sees it. +@runtime_checkable +class HierarchicalLiveWeightArenaProvider(Protocol): + """Provider boundary implemented by TEKit or a Fabric test allocator.""" + + def build_hierarchical_live_arena( + self, + *, + world: int, + rank: int, + home_count: int, + helper_count: int, + bundle: BundleLayout, + device: int, + comm: object, + ) -> HierarchicalLiveWeightArena: + """Construct the external live bank and all aligned-group aliases.""" + + +def _stream_handle(stream: object) -> int: + if isinstance(stream, bool): + raise TypeError("stream must be a CUDA stream handle") + handle = getattr(stream, "cuda_stream", stream) + value = int(handle) + if value <= 0: + raise ValueError("hierarchical copy requires a positive non-default stream") + return value + + +def _event_handle(event: object | None) -> int: + if event is None: + return 0 + if isinstance(event, bool): + raise TypeError("event must be a CUDA event handle") + return int(getattr(event, "cuda_event", event)) + + +@dataclass(frozen=True) +class HierarchicalCopyTicket: + generation: int + terminal_flags: torch.Tensor = field(repr=False) + command_count: int + + +class HierarchicalSamiWeightBroadcast: + """One-rank HALO-Q copy endpoint with a single C hot-path context.""" + + partition = "bundle" + remote_visibility_release = True + live_bank_count = 1 + + def __init__( + self, + *, + comm: object, + device: int, + bundle: BundleLayout, + helper_count: int, + global_expert_count: int, + arena_provider: HierarchicalLiveWeightArenaProvider, + tma_sm_count: int | None = None, + tma_warps: int | None = None, + tma_route: str = "plan", + tma_source_load_factor: float = 1.0, + ) -> None: + for name, value in (("tma_sm_count", tma_sm_count), ("tma_warps", tma_warps)): + if value is not None and (type(value) is not int or value <= 0): + raise ValueError(f"{name} must be a positive exact int or None") + if tma_sm_count is not None and tma_sm_count > 8: + raise ValueError("tma_sm_count must be at most 8") + if tma_warps is not None and tma_warps > 32: + raise ValueError("tma_warps must be at most 32; device shared-memory limits also apply") + if tma_route not in ("plan", "source", "scatter"): + raise ValueError("tma_route must be 'plan', 'source', or 'scatter'") + if (type(tma_source_load_factor) not in (int, float) or + not 0.01 <= tma_source_load_factor < 2.0 or + abs(100 * tma_source_load_factor - round(100 * tma_source_load_factor)) > 1e-9): + raise ValueError("tma_source_load_factor must be 0.01..1.99 in steps of 0.01") + self.tma_route = tma_route + self.tma_source_load_factor = round(100 * tma_source_load_factor) / 100 + world = getattr(comm, "world", None) + rank = getattr(comm, "rank", None) + if type(world) is not int: + raise ValueError("comm.world must be an exact int") + hierarchy_group_sizes(world) + if type(rank) is not int or not 0 <= rank < world: + raise ValueError("comm.rank must be in [0, world)") + if type(device) is not int or device < 0: + raise ValueError("device must be a non-negative exact int") + if ( + type(global_expert_count) is not int + or global_expert_count <= 0 + or global_expert_count % world + ): + raise ValueError("global_expert_count must be positive and EP-divisible") + if type(helper_count) is not int or helper_count <= 0: + raise ValueError("helper_count must be a positive exact int") + if global_expert_count + world * helper_count > (1 << 31) - 1: + raise ValueError("physical slots exceed public int32 indexing") + if 4 + 7 * ((helper_count + 3) // 4 * 4) > (1 << 31) - 1: + raise ValueError("PlanChannel word count exceeds native int32 indexing") + if not isinstance(arena_provider, HierarchicalLiveWeightArenaProvider): + raise TypeError("arena_provider lacks the hierarchical arena protocol") + self.comm = comm + self.device = device + self.bundle = bundle + self.helper_count = helper_count + self.global_expert_count = global_expert_count + self.home_count = global_expert_count // world + self.home_begin = rank * self.home_count + self.provider = arena_provider + self.cuda = cuda + torch.cuda.set_device(device) + # Fail launch-geometry checks before allocating collective Fabric resources. + device_sms = torch.cuda.get_device_properties(device).multi_processor_count + if tma_sm_count is not None and tma_sm_count > device_sms: + raise ValueError(f"tma_sm_count must be at most {device_sms} on this GPU") + self._validate_tma_warp_capacity(tma_warps) + # Unlike CTA/warp counts, the ownership rule must agree across ranks. + # This is one cold-path collective, outside every copy submission. + factors = comm.allgather(self.tma_source_load_factor) + if any(value != self.tma_source_load_factor for value in factors): + raise ValueError("tma_source_load_factor must agree across all ranks") + + live = arena_provider.build_hierarchical_live_arena( + world=world, + rank=rank, + home_count=self.home_count, + helper_count=helper_count, + bundle=bundle, + device=device, + comm=comm, + ) + if not isinstance(live, HierarchicalLiveWeightArena): + raise TypeError("provider returned a non-hierarchical live arena") + self.live_arena = live + self.arena = live.bind( + world=world, + rank=rank, + home_count=self.home_count, + helper_count=helper_count, + bundle=bundle, + device=device, + ) + self.live_plane_tensors = self.arena.base.local_plane_views + self.terminal_flags = self.live_arena.base.terminals.local_view + self.flags_local = self.arena.base.terminal_local_ptr + self.flags_mc = self.arena.base.terminal_mc_ptr + 8 * rank + check_cuda( + cuda.cuMemsetD8( + cuda.CUdeviceptr(self.flags_local), 0, 8 * world + ), + "cuMemsetD8(hierarchical terminals)", + ) + torch.cuda.synchronize(device) + comm.barrier() + + # GPU-direct consumes the scheduler's device storage through the + # statically linked THOP; no mapped channel or runtime compilation exists. + plan_stride = (helper_count + 3) // 4 * 4 + plan_words = 4 + 7 * plan_stride + planes = len(bundle.planes) + ctx = torch.ops.trtllm.moe_rebalance_tma_create( + helper_count * planes, tma_sm_count or 0, tma_warps or 0 + ) + try: + target_count = sum(world // size for size in self.arena.group_sizes[1:]) + torch.ops.trtllm.moe_rebalance_tma_configure_gpu_plan( + ctx, world, rank, helper_count, planes, + global_expert_count, self.home_count, plan_stride, + len(self.arena.group_sizes), target_count, 6, plan_words, 0, + {"plan": 0, "source": 1, "scatter": 2}[tma_route], + round(100 * self.tma_source_load_factor), + list(self.arena.group_sizes), + list(self.arena.base.scatter_source_ptrs), + list(self.arena.group_destination_ptrs), [], + [plane.nbytes for plane in bundle.planes], + ) + config_names = ( + "abi_version", "device", "device_sm_count", "sms", "warps", + "threads_per_cta", "slots_per_warp", "bank0_slots_per_warp", + "bank1_slots_per_warp", "slice_bytes", "dynamic_shared_bytes", + "device_optin_shared_bytes", "device_shared_bytes_per_sm", + "max_active_ctas_per_sm", "compute_major", "compute_minor", + "total_slots", "max_slots_per_warp", "extra_slot_warps", + "max_warps", "max_segments", + ) + self.tma_config = dict(zip( + config_names, torch.ops.trtllm.moe_rebalance_tma_config(ctx) + )) + if int(torch.ops.trtllm.moe_rebalance_tma_current_gen(ctx)) != 1: + raise RuntimeError("hierarchical SAMI context is not fresh") + except BaseException: + torch.ops.trtllm.moe_rebalance_tma_destroy(ctx) + raise + self.ctx = ctx + + self._lock = threading.Lock() + self._scheduler: object | None = None + self._outputs: object | None = None + self._stream = 0 + self._stream_owner: object | None = None + self._generation = 1 + self._reuse_authority: object | None = None + self._active_generation = 0 + self._local_release_generation = 0 + self._collective_safe_generation = 0 + self._handoff_aborted = False + + def _validate_tma_warp_capacity(self, warps: int | None) -> None: + # Reject impossible device geometry before building a collective arena. + # Native creation additionally checks the compiled kernel's own static + # shared memory, thread limit and occupancy. + device = check_cuda(self.cuda.cuDeviceGet(self.device), "cuDeviceGet(TMA geometry)") + shared_bytes = int(check_cuda(self.cuda.cuDeviceGetAttribute( + self.cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, + device, + ), "cuDeviceGetAttribute(TMA opt-in shared memory)")) + max_threads = int(check_cuda(self.cuda.cuDeviceGetAttribute( + self.cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK, + device, + ), "cuDeviceGetAttribute(TMA block threads)")) + slot_bytes = 8208 + 32 # Logical 8 KiB plus aligned-load padding and SlotControl. + slots = max(0, shared_bytes // slot_bytes) + while slots and ((slots * slot_bytes + 127) & ~127) > shared_bytes: + slots -= 1 + max_warps = min(32, slots // 2, max_threads // 64) + if (warps or 7) > max_warps: + raise ValueError( + f"tma_warps must be at most {max_warps} on this GPU " + "(at least two 8 KiB slices per worker pair)" + ) + + def _validate_stream_context(self, handle: int) -> None: + get_current = getattr(self.cuda, "cuCtxGetCurrent", None) + get_device = getattr(self.cuda, "cuCtxGetDevice", None) + get_stream_context = getattr(self.cuda, "cuStreamGetCtx", None) + if not callable(get_current) or not callable(get_device): + raise RuntimeError("CUDA driver binding cannot validate the current context") + current_context = check_cuda(get_current(), "cuCtxGetCurrent") + if current_context is None: + raise RuntimeError("no current CUDA context for hierarchical scheduler") + current_device = int(check_cuda(get_device(), "cuCtxGetDevice")) + if current_device != self.device: + raise ValueError("current CUDA context and SAMI devices differ") + if handle: + if not callable(get_stream_context): + raise RuntimeError("CUDA driver binding cannot validate stream context") + stream_context = check_cuda( + get_stream_context(handle), "cuStreamGetCtx(hierarchical scheduler)" + ) + if stream_context != current_context: + raise ValueError("scheduler stream is from a different CUDA context") + + def bind_scheduler(self, scheduler: object, stream: object) -> None: + """Bind the four-field CUDA scheduler and one static stream for life.""" + + with self._lock: + if self._handoff_aborted: + raise RuntimeError( + "hierarchical handoff aborted; quiesced group rebuild required" + ) + if self._scheduler is not None: + raise RuntimeError("hierarchical scheduler binding is single-assignment") + config = getattr(scheduler, "cfg", None) + outputs = getattr(scheduler, "outputs", None) + if config is None or outputs is None: + raise TypeError("scheduler lacks cfg/outputs") + + expected = ( + ("ep_size", int(self.comm.world)), + ("logical_expert_count", self.global_expert_count), + ("extra_slots_per_rank", self.helper_count), + ("local_rank", int(self.comm.rank)), + ) + for name, value in expected: + observed = getattr(config, name, None) + if type(observed) is not int or observed != value: + raise ValueError(f"scheduler {name} is incompatible with SAMI") + for name in ("max_tokens_per_rank", "topk"): + observed = getattr(config, name, None) + if type(observed) is not int or observed <= 0: + raise ValueError(f"scheduler {name} must be a positive exact int") + fields = ( + "physical_slot_ids", + "hot_expert_ids", + "hot_expert_group_level", + "hot_expert_source_ranks", + ) + tensors = tuple(getattr(outputs, name, None) for name in fields) + if any(not isinstance(tensor, torch.Tensor) for tensor in tensors): + raise TypeError("scheduler does not expose the four-tensor HALO ABI") + expected_shapes = ( + (config.max_tokens_per_rank, config.topk), + (self.helper_count,), + (self.helper_count,), + (self.helper_count,), + ) + for name, tensor, expected_shape in zip(fields, tensors, expected_shapes): + if ( + tensor.dtype != torch.int32 + or not tensor.is_cuda + or not tensor.is_contiguous() + ): + raise ValueError( + f"{name} must be contiguous CUDA int32 storage" + ) + if tuple(tensor.shape) != expected_shape: + raise ValueError(f"{name} must have shape {expected_shape}") + + scheduler_device = getattr(scheduler, "device", None) + try: + scheduler_device = torch.device(scheduler_device) + except (TypeError, RuntimeError) as exc: + raise TypeError("scheduler must expose one explicit CUDA device") from exc + if scheduler_device.type != "cuda" or scheduler_device.index is None: + raise ValueError("scheduler must expose one explicit CUDA device") + if scheduler_device.index != self.device: + raise ValueError("scheduler and SAMI devices differ") + if any(tensor.device != scheduler_device for tensor in tensors): + raise ValueError("HALO outputs must share the scheduler CUDA device") + + handle = _stream_handle(stream) + stream_device = getattr(stream, "device", None) + if stream_device is not None: + try: + stream_device = torch.device(stream_device) + except (TypeError, RuntimeError) as exc: + raise TypeError("scheduler stream exposes an invalid device") from exc + if stream_device != scheduler_device: + raise ValueError("scheduler stream and output devices differ") + bind = getattr(scheduler, "bind_gpu_direct", None) + release = getattr(scheduler, "release_plan_channel", None) + if not callable(bind) or not callable(release): + raise TypeError("scheduler lacks the copy handoff lifecycle") + workspace = getattr(scheduler, "plan_workspace", None) + status = getattr(scheduler, "status", None) + grid_sync = getattr(scheduler, "grid_sync", None) + capacity = getattr(scheduler, "max_broadcasts", None) + if type(capacity) is not int or capacity <= 0: + raise ValueError("GPU-direct requires the scheduler plan capacity") + if (not isinstance(workspace, torch.Tensor) or + workspace.device != scheduler_device or workspace.dtype != torch.int32 or + not workspace.is_contiguous() or tuple(workspace.shape) != (2 + 6 * capacity,)): + raise ValueError("GPU-direct requires the original CUDA scheduler plan workspace") + if (not isinstance(status, torch.Tensor) or status.device != scheduler_device or + status.dtype != torch.int32 or not status.is_contiguous() or status.numel() < 6): + raise ValueError("GPU-direct requires the CUDA scheduler status workspace") + if (not isinstance(grid_sync, torch.Tensor) or grid_sync.device != scheduler_device or + grid_sync.dtype != torch.int32 or not grid_sync.is_contiguous() or grid_sync.numel() < 2): + raise ValueError("GPU-direct requires the CUDA scheduler epoch workspace") + self._validate_stream_context(handle) + try: + bind(handle) + torch.ops.trtllm.moe_rebalance_tma_bind_gpu_direct( + self.ctx, *tensors[1:], workspace, status, grid_sync, capacity + ) + except BaseException: + self._handoff_aborted = True + raise + self._scheduler = scheduler + self._outputs = outputs + self._stream = handle + self._stream_owner = stream + + def _validate_reuse_authority(self, authority: object) -> None: + if authority is None: + raise TypeError("generation reuse authority cannot be None") + if getattr(authority, "collective_safe", False) is not True: + raise ValueError("generation reuse authority lacks collective safety") + if getattr(authority, "producer_reuse_guarded", False) is not True: + raise ValueError("generation reuse authority does not guard producer reuse") + if getattr(authority, "live_bank_count", None) != self.live_bank_count: + raise ValueError("generation reuse authority must target one live bank") + copy_module = getattr(authority, "bound_copy_module", None) + if copy_module is not self and getattr(copy_module, "backend", None) is not self: + raise ValueError( + "generation reuse authority is not bound to this hierarchical endpoint" + ) + if copy_module is not self: + config = getattr(copy_module, "config", None) + expected = { + "ep_size": int(self.comm.world), + "logical_expert_count": self.global_expert_count, + "extra_slots_per_rank": self.helper_count, + "local_rank": int(self.comm.rank), + } + if config is None or any( + type(getattr(config, name, None)) is not int + or getattr(config, name) != value + for name, value in expected.items() + ): + raise ValueError("generation reuse authority geometry is inconsistent") + + def bind_generation_reuse_authority(self, authority: object) -> None: + """Bind the exact collective single-bank lease provider once.""" + + self._validate_reuse_authority(authority) + with self._lock: + if self._handoff_aborted: + raise RuntimeError( + "hierarchical handoff aborted; quiesced group rebuild required" + ) + if self._reuse_authority is not None: + raise RuntimeError("generation reuse authority is single-assignment") + if self._active_generation: + raise RuntimeError("generation reuse authority must bind before submit") + self._reuse_authority = authority + + def release_generation_after( + self, generation: int, completion_event: object + ) -> None: + """Order the copy stream after this rank's live-bank consumer event. + + This establishes only local stream ordering. The bound reuse authority + must separately attest its bounded all-rank barrier before another + generation may overwrite the single live bank. + """ + + if type(generation) is not int or generation <= 0: + raise ValueError("generation must be a positive exact int") + event_handle = _event_handle(completion_event) + if event_handle <= 0: + raise ValueError("completion_event must expose a positive CUDA event") + with self._lock: + if self._handoff_aborted or self._scheduler is None: + raise RuntimeError("hierarchical scheduler handoff is not operational") + if generation != self._active_generation: + raise ValueError("generation is not the active live bank") + if self._local_release_generation >= generation: + raise RuntimeError("generation already has a local release event") + check_cuda( + self.cuda.cuStreamWaitEvent(self._stream, event_handle, 0), + "cuStreamWaitEvent(live-bank consumer)", + ) + self._local_release_generation = generation + + def mark_collective_reuse_safe( + self, authority: object, generation: int + ) -> None: + """Accept the bound provider's all-rank proof for one generation.""" + + if type(generation) is not int or generation <= 0: + raise ValueError("generation must be a positive exact int") + with self._lock: + if self._handoff_aborted or self._scheduler is None: + raise RuntimeError("hierarchical scheduler handoff is not operational") + if authority is not self._reuse_authority: + raise PermissionError("generation reuse authority identity mismatch") + self._validate_reuse_authority(authority) + if generation != self._active_generation: + raise ValueError("generation is not the active live bank") + if self._local_release_generation != generation: + raise RuntimeError("local consumer release is not registered") + if self._collective_safe_generation >= generation: + raise RuntimeError("generation is already collectively reusable") + self._collective_safe_generation = generation + + def _require_live_bank_reusable(self) -> None: + if self._reuse_authority is None: + raise RuntimeError("single live bank has no bound collective reuse authority") + self._validate_reuse_authority(self._reuse_authority) + if ( + self._active_generation + and self._collective_safe_generation < self._active_generation + ): + raise RuntimeError( + "single live bank is still leased by the previous generation" + ) + + def submit( + self, + outputs: object, + *, + start_event: object | None = None, + end_event: object | None = None, + payload_end_event: object | None = None, + ) -> HierarchicalCopyTicket: + """Enqueue the fused GPU-direct TMA copy and READY publication.""" + + with self._lock: + if self._handoff_aborted: + raise RuntimeError( + "hierarchical handoff aborted; quiesced group rebuild required" + ) + if outputs is not self._outputs or self._scheduler is None: + raise RuntimeError("outputs are not from the bound CUDA scheduler") + self._require_live_bank_reusable() + try: + with torch.cuda.device(self.device), torch.cuda.stream( + self._stream_owner + ): + if start_event is not None: + start_event.record() + commands = int( + torch.ops.trtllm.moe_rebalance_tma_submit_gpu_direct( + self.ctx, self.flags_mc + ) + ) + if payload_end_event is not None: + payload_end_event.record() + if end_event is not None: + end_event.record() + self._active_generation = self._generation + release = getattr(self._scheduler, "release_plan_channel") + release() + except BaseException: + self._handoff_aborted = True + raise + ticket = HierarchicalCopyTicket( + generation=self._generation, + terminal_flags=self.terminal_flags, + command_count=commands + 1, + ) + self._generation += 1 + return ticket + + def close(self) -> None: + """Release the native copy context after all device work is drained.""" + with self._lock, torch.cuda.device(self.device): + if self.ctx is None: + return + torch.ops.trtllm.moe_rebalance_tma_destroy(self.ctx) + self.ctx = None + self._scheduler = None + self._outputs = None + self._stream_owner = None + + def last_mode_counts(self) -> tuple[int, int, int]: + """Return diagnostic ``(direct_slots, scatter_slots, commands)``. + + GPU-direct mode explicitly waits for the copy and reads GPU counters. + This opt-in diagnostic must stay outside the submission/timing path. + """ + + result = torch.ops.trtllm.moe_rebalance_tma_result(self.ctx) + direct, scatter, commands = result[1], result[2], result[7] + return int(direct), int(scatter), int(commands) + + +__all__ = [ + "BoundHierarchicalLiveWeightArena", + "HierarchicalCopyTicket", + "HierarchicalLiveWeightArena", + "HierarchicalLiveWeightArenaProvider", + "HierarchicalSamiWeightBroadcast", + "hierarchy_group_sizes", +] diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_shared_fc12.py b/tensorrt_llm/_torch/moe/kernels/megamoe_shared_fc12.py new file mode 100644 index 000000000000..c197354ae6f6 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_shared_fc12.py @@ -0,0 +1,652 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Single-expert use of MegaMoE's complete fused FC1/activation/FC2 kernel.""" + +from itertools import product +from threading import Condition, RLock + +import torch + +from ..autotuner import AutoTuner, DynamicTensorSpec, TunableRunner, TuningConfig + +__all__ = [ + "release_shared_fc12_cache", + "run_shared_fc12", + "shared_fc12_cache_info", + "shared_fc12_candidate_tactics", + "shared_fc12_tactic_descriptor", +] + +_SHARED_FC12_OP = "trtllm::megamoe_shared_fc12" +_CATALOG_VERSION = 1 +_SharedFc12Tactic = tuple[int, int, int, str, bool] +_CANDIDATES = tuple( + product((1, 2), (64, 128, 256), (128, 256), ("grid_stride", "atomic_counter"), (False, True)) +) + + +def shared_fc12_candidate_tactics() -> tuple[_SharedFc12Tactic, ...]: + """Finite common search space; larger clusters/K and extra TMA stages are not searched.""" + return _CANDIDATES + + +def shared_fc12_tactic_descriptor(tactic: _SharedFc12Tactic) -> dict: + tactic = tuple(tactic) + if tactic not in _CANDIDATES: + raise ValueError(f"Unsupported shared FC12 tactic: {tactic}") + ctas, tile_n, tile_k, work_id_mode, fc2_use_bulk = tactic + return { + "tactic": list(tactic), + "mma_tiler_mnk": [128 * ctas, tile_n, tile_k], + "cluster_shape_mn": [ctas, 1], + "use_2cta_instrs": ctas == 2, + "work_id_mode": work_id_mode, + "fc2_use_bulk": fc2_use_bulk, + "fc2_tma_stages": 1, + "fc1_epi_flag_batch": 1, + "fc2_epi_flag_batch": 1, + } + + +def _token_buckets(max_tokens: int) -> tuple[int, ...]: + if max_tokens <= 0: + raise ValueError("shared FC12 capacity must be positive") + return tuple( + sorted( + {n for n in (1, 128, 256, 512, 1024, 2048, 4096, 8192) if n < max_tokens} | {max_tokens} + ) + ) + + +def _make_kernel( + device_index, hidden_size, intermediate_size, max_tokens, sm_count, swiglu_limit, tactic +): + import cutlass.utils as utils + from cutlass.cute.nvgpu import OperandMajorMode + + from .cutedsl_megamoe.api import ImplDesc, ProblemDesc + from .cutedsl_megamoe.kernel_src.blackwell.inference.mega.block_scaled_swap_ab_fc12_kernel import ( + BlockScaledSwapAbFc12Kernel, + ) + from .cutedsl_megamoe.quant_def import CombineFormat + + properties = torch.cuda.get_device_properties(device_index) + if not 0 < sm_count <= properties.multi_processor_count: + raise ValueError("shared FC12 sm_count must be within the device SM count") + descriptor = shared_fc12_tactic_descriptor(tactic) + cluster_size = descriptor["cluster_shape_mn"][0] + hardware_clusters = utils.HardwareInfo(device_index).get_max_active_clusters(cluster_size) + launch_clusters = min(sm_count // cluster_size, hardware_clusters) + if launch_clusters <= 0: + raise ValueError("shared FC12 requires at least one active cluster") + problem = ProblemDesc( + { + "expert_count": 1, + "intermediate_gateup_size": 2 * intermediate_size, + "hidden_size": hidden_size, + "quant_kind": "mxfp8_e4m3", + "a_major_mode": OperandMajorMode.K, + "b_major_mode": OperandMajorMode.K, + "combine_format": CombineFormat.parse("bf16"), + "gate_up_clamp": swiglu_limit, + } + ) + implementation = ImplDesc( + { + **{key: value for key, value in descriptor.items() if key != "tactic"}, + "mma_tiler_mnk": tuple(descriptor["mma_tiler_mnk"]), + "cluster_shape_mn": tuple(descriptor["cluster_shape_mn"]), + "group_hint": launch_clusters, + "token_padding_block": 64, + "sf_padding_block": 128, + "max_tokens": max_tokens, + "launch_cluster_count": launch_clusters, + } + ) + kernel = BlockScaledSwapAbFc12Kernel(problem, implementation) + return kernel, implementation, properties, hardware_clusters, launch_clusters + + +class _SharedFc12Runner: + def __init__( + self, + device_index: int, + hidden_size: int, + intermediate_size: int, + max_tokens: int, + sm_count: int, + swiglu_limit: float | None, + stream_handle: int, + *, + tactic: _SharedFc12Tactic | None = None, + ) -> None: + import cuda.bindings.driver as cuda + import cutlass + import cutlass.cute as cute + from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream, make_ptr + from cutlass.cute.typing import AddressSpace + + device = torch.device("cuda", device_index) + properties = torch.cuda.get_device_properties(device) + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("Create the shared FC12 runner before CUDA graph capture") + self.hidden_size = hidden_size + self.max_tokens = max_tokens + self.reserved_sms = properties.multi_processor_count - sm_count + self.stream_handle = stream_handle + if tactic is None: + tactic = ( + (1, 256, 128, "atomic_counter", True) + if self.reserved_sms > 0 + else (2, 128, 128, "grid_stride", True) + ) + self.tactic = tuple(tactic) + kernel, implementation, properties, hardware_clusters, launch_clusters = _make_kernel( + device_index, + hidden_size, + intermediate_size, + max_tokens, + sm_count, + swiglu_limit, + self.tactic, + ) + local_bytes, shared_bytes = kernel.get_workspace_sizes() + if shared_bytes: + raise RuntimeError("shared FC12 must not allocate a communication workspace") + self.workspace = torch.empty(local_bytes, dtype=torch.uint8, device=device) + local_zero_bytes, shared_zero_bytes = kernel.require_zero_workspace_leading_bytes + if shared_zero_bytes: + raise RuntimeError("shared FC12 must not initialize communication state") + self.workspace_pointer = self.workspace.data_ptr() + self.workspace_reset_bytes = local_zero_bytes + self._memset_async = cuda.cuMemsetD8Async + self._cuda_success = cuda.CUresult.CUDA_SUCCESS + self._cuda_stream = cuda.CUstream(stream_handle) + # A slice supplies the current E=1 token count without a per-call GPU write. + self.token_counts = torch.arange(max_tokens + 1, dtype=torch.int32, device=device) + + token_rows = cute.sym_int64(divisibility=1) + scale_rows = cute.sym_int64(divisibility=128) + padded_gateup = (2 * intermediate_size + 127) // 128 * 128 + fc1_sf_columns = padded_gateup * (hidden_size // 32) + fc2_sf_columns = hidden_size * (intermediate_size // 32) + fake_arguments = { + "activation": make_fake_compact_tensor( + cutlass.Float8E4M3FN, + (token_rows, hidden_size), + stride_order=(1, 0), + assumed_align=16, + ), + "fc1_weight": make_fake_compact_tensor( + cutlass.Float8E4M3FN, + (1, hidden_size, 2 * intermediate_size), + stride_order=(2, 0, 1), + assumed_align=16, + ), + "activation_sf": make_fake_compact_tensor( + cutlass.Float8E8M0FNU, + (scale_rows, hidden_size // 32), + stride_order=(1, 0), + assumed_align=16, + ), + "fc1_weight_sf": make_fake_compact_tensor( + cutlass.Float8E8M0FNU, (1, fc1_sf_columns), stride_order=(1, 0), assumed_align=16 + ), + "fc2_weight": make_fake_compact_tensor( + cutlass.Float8E4M3FN, + (1, intermediate_size, hidden_size), + stride_order=(2, 0, 1), + assumed_align=16, + ), + "fc2_weight_sf": make_fake_compact_tensor( + cutlass.Float8E8M0FNU, (1, fc2_sf_columns), stride_order=(1, 0), assumed_align=16 + ), + "fc2_output": make_fake_compact_tensor( + cutlass.BFloat16, + (token_rows, 1, hidden_size), + stride_order=(2, 1, 0), + assumed_align=16, + ), + "local_workspace": make_ptr(cutlass.Uint8, 0, AddressSpace.gmem, assumed_align=128), + "max_active_clusters": hardware_clusters, + "stream": make_fake_stream(), + "expert_token_sizes": make_fake_compact_tensor( + cutlass.Int32, (1,), stride_order=(0,), assumed_align=4 + ), + "expert_token_prefix_sum": None, + "topk_scores": None, + "fc1_alpha": None, + "fc2_alpha": None, + "fc1_norm_const": None, + } + self.compiled = cute.compile[cute.EnableTVMFFI(True)](kernel, **fake_arguments) + + def __call__( + self, + activation: torch.Tensor, + fc1_weight: torch.Tensor, + fc1_weight_sf: torch.Tensor, + fc2_weight: torch.Tensor, + fc2_weight_sf: torch.Tensor, + ) -> torch.Tensor: + token_rows = activation.shape[0] + quantized, scales = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0.sm_budget( + activation, reserved_sms=self.reserved_sms + ) + output = torch.empty( + (token_rows, 1, self.hidden_size), dtype=torch.bfloat16, device=activation.device + ) + # The standalone FC12 readiness counters require reset before every invocation. + (status,) = self._memset_async( + self.workspace_pointer, 0, self.workspace_reset_bytes, self._cuda_stream + ) + if status != self._cuda_success: + raise RuntimeError(f"shared FC12 asynchronous counter reset failed: {status}") + self.compiled( + activation=quantized, + fc1_weight=fc1_weight.transpose(1, 2), + activation_sf=scales.view(torch.float8_e8m0fnu).view(-1, self.hidden_size // 32), + fc1_weight_sf=fc1_weight_sf.view(torch.float8_e8m0fnu), + fc2_weight=fc2_weight.transpose(1, 2), + fc2_weight_sf=fc2_weight_sf.view(torch.float8_e8m0fnu), + fc2_output=output, + local_workspace=self.workspace_pointer, + stream=self.stream_handle, + expert_token_sizes=self.token_counts[token_rows : token_rows + 1], + expert_token_prefix_sum=None, + topk_scores=None, + fc1_alpha=None, + fc2_alpha=None, + fc1_norm_const=None, + ) + return output[:, 0, :] + + +_SHARED_FC12_CACHE_LOCK = RLock() +_SHARED_FC12_CACHE_CONDITION = Condition(_SHARED_FC12_CACHE_LOCK) +_ACTIVE_FC12_CALLS: dict[tuple[int, int], int] = {} +_RELEASING_FC12_DEVICES: set[int] = set() +_RELEASING_FC12_STREAMS: set[tuple[int, int]] = set() +_RUNNER_CACHE: dict[tuple, _SharedFc12Runner] = {} +_RUNNER_CACHE_HITS = 0 +_RUNNER_CACHE_MISSES = 0 + + +def _begin_shared_fc12_release( + device_index: int, stream_handle: int | None +) -> tuple[int, int] | None: + stream_key = (device_index, stream_handle) if stream_handle is not None else None + with _SHARED_FC12_CACHE_CONDITION: + while ( + device_index in _RELEASING_FC12_DEVICES + or ( + stream_key is None + and any(device == device_index for device, _ in _RELEASING_FC12_STREAMS) + ) + or (stream_key is not None and stream_key in _RELEASING_FC12_STREAMS) + ): + _SHARED_FC12_CACHE_CONDITION.wait() + if stream_key is None: + _RELEASING_FC12_DEVICES.add(device_index) + else: + _RELEASING_FC12_STREAMS.add(stream_key) + while any( + count and device == device_index and (stream_handle is None or stream == stream_handle) + for (device, stream), count in _ACTIVE_FC12_CALLS.items() + ): + _SHARED_FC12_CACHE_CONDITION.wait() + return stream_key + + +def _end_shared_fc12_release(device_index: int, stream_key: tuple[int, int] | None) -> None: + with _SHARED_FC12_CACHE_CONDITION: + if stream_key is None: + _RELEASING_FC12_DEVICES.discard(device_index) + else: + _RELEASING_FC12_STREAMS.discard(stream_key) + _SHARED_FC12_CACHE_CONDITION.notify_all() + + +def _begin_shared_fc12_call(device_index: int, stream_handle: int) -> tuple[int, int]: + key = (int(device_index), int(stream_handle)) + with _SHARED_FC12_CACHE_CONDITION: + while key[0] in _RELEASING_FC12_DEVICES or key in _RELEASING_FC12_STREAMS: + _SHARED_FC12_CACHE_CONDITION.wait() + _ACTIVE_FC12_CALLS[key] = _ACTIVE_FC12_CALLS.get(key, 0) + 1 + return key + + +def _end_shared_fc12_call(key: tuple[int, int]) -> None: + with _SHARED_FC12_CACHE_CONDITION: + remaining = _ACTIVE_FC12_CALLS[key] - 1 + if remaining: + _ACTIVE_FC12_CALLS[key] = remaining + else: + del _ACTIVE_FC12_CALLS[key] + _SHARED_FC12_CACHE_CONDITION.notify_all() + + +def _get_runner( + device_index: int, + hidden_size: int, + intermediate_size: int, + max_tokens: int, + sm_count: int, + swiglu_limit: float | None, + stream_handle: int, + tactic: _SharedFc12Tactic | None = None, +) -> _SharedFc12Runner: + global _RUNNER_CACHE_HITS, _RUNNER_CACHE_MISSES + key = ( + device_index, + hidden_size, + intermediate_size, + max_tokens, + sm_count, + swiglu_limit, + stream_handle, + tactic, + ) + with _SHARED_FC12_CACHE_LOCK: + runner = _RUNNER_CACHE.get(key) + if runner is not None: + _RUNNER_CACHE_HITS += 1 + return runner + _RUNNER_CACHE_MISSES += 1 + with torch.cuda.device(device_index): + runner = _SharedFc12Runner( + device_index, + hidden_size, + intermediate_size, + max_tokens, + sm_count, + swiglu_limit, + stream_handle, + tactic=tactic, + ) + _RUNNER_CACHE[key] = runner + return runner + + +def shared_fc12_cache_info() -> dict[str, int | None]: + """Return process-cache counters for diagnostics.""" + with _SHARED_FC12_CACHE_LOCK: + return { + "hits": _RUNNER_CACHE_HITS, + "misses": _RUNNER_CACHE_MISSES, + "maxsize": None, + "currsize": len(_RUNNER_CACHE), + } + + +class SharedFc12TunableRunner(TunableRunner): + # AutoTuner uses this dictionary to prime persisted winners during model warmup. + kernel_cache: dict[tuple, _SharedFc12Runner] = {} + + def __init__( + self, + device_index, + hidden_size, + intermediate_size, + max_tokens, + sm_count, + swiglu_limit, + stream_handle, + input_signature, + ): + properties = torch.cuda.get_device_properties(device_index) + self.device_index = device_index + self.sm_count = sm_count + self.reserved_sms = properties.multi_processor_count - sm_count + self.max_tokens = max_tokens + self.stream_handle = stream_handle + self._runner_args = ( + device_index, + hidden_size, + intermediate_size, + max_tokens, + sm_count, + swiglu_limit, + stream_handle, + ) + self._unique_id = ( + _CATALOG_VERSION, + device_index, + properties.major, + properties.minor, + properties.multi_processor_count, + sm_count, + hidden_size, + intermediate_size, + max_tokens, + swiglu_limit, + "K128_input_K32_intermediate_R128c4_gateup16", + input_signature, + ) + self.buckets = _token_buckets(max_tokens) + self.tuning_config = TuningConfig( + dynamic_tensor_specs=( + DynamicTensorSpec( + input_idx=0, + dim_idx=0, + gen_tuning_buckets=self.buckets, + map_to_tuning_buckets=self.token_bucket, + ), + ), + tune_max_num_tokens=max_tokens, + use_cold_l2_cache=True, + use_cuda_graph=False, + ) + self._valid_tactics = None + self._compiled_runners = {} + self._compile_lock = RLock() + + def unique_id(self): + return self._unique_id + + def token_bucket(self, tokens: int) -> int: + if not 0 < tokens <= self.max_tokens: + raise ValueError("shared FC12 token count exceeds its configured capacity") + return next(bucket for bucket in self.buckets if tokens <= bucket) + + def get_valid_tactics(self, inputs, profile, **kwargs): + if self._valid_tactics is None: + valid = [] + with torch.cuda.device(self.device_index): + for tactic in shared_fc12_candidate_tactics(): + try: + _make_kernel(*self._runner_args[:6], tactic) + except (ValueError, AssertionError): + continue + else: + valid.append(tactic) + self._valid_tactics = tuple(valid) + return list(self._valid_tactics) + + def _cached_winner_tactics(self, tuner, inputs): + """Return every profiled bucket winner, or None if any is missing.""" + if tuner.skip_dynamic_tuning_buckets: + return None + winners = set() + profiles = tuner._optimization_profiles(self.tuning_config, inputs) + if not profiles: + return None + for profile in profiles: + hit, runner_id, tactic, _ = tuner.profiling_cache.search_cache( + _SHARED_FC12_OP, + [self], + profile.get_opt_shapes(), + self.tuning_config, + apply_map_to_tuning_buckets=False, + ) + if not hit or runner_id != 0 or tactic == -1: + return None + tactic = tuple(tactic) + if tactic not in _CANDIDATES: + return None + winners.add(tactic) + return winners + + def _retain_cached_winners(self, winners) -> None: + """Drop profiled loser workspaces after all bucket winners are known.""" + winners = {tuple(tactic) for tactic in winners} + if not winners or not winners.issubset(_CANDIDATES): + return + with self._compile_lock: + stale = set(self._compiled_runners) - winners + if not stale: + return + # Every cached runner is bound to this stream. Profiling has + # finished, but synchronize before releasing its workspace. + torch.cuda.ExternalStream(self.stream_handle, device=self.device_index).synchronize() + with _SHARED_FC12_CACHE_LOCK: + for tactic in stale: + self._compiled_runners.pop(tactic, None) + key = (*self._runner_args, tactic) + self.kernel_cache.pop(key, None) + _RUNNER_CACHE.pop(key, None) + # Return the released loser workspaces to CUDA before the routed + # MegaMoE autotuner allocates its symmetric profiling scratch. + with torch.cuda.device(self.device_index): + torch.cuda.empty_cache() + + def forward(self, inputs, *, tactic=-1, do_preparation=False, **kwargs): + if do_preparation: + return None + if tactic == -1: + if AutoTuner.get().is_tuning_mode: + raise RuntimeError("shared FC12 profiling did not select a valid tactic") + tactic = ( + (1, 256, 128, "atomic_counter", True) + if self.reserved_sms > 0 + else (2, 128, 128, "grid_stride", True) + ) + tactic = tuple(tactic) + with self._compile_lock: + runner = self._compiled_runners.get(tactic) + if runner is None: + runner = _get_runner(*self._runner_args, tactic) + self._compiled_runners[tactic] = runner + self.kernel_cache[(*self._runner_args, tactic)] = runner + return runner(*inputs) + + +_tunable_runners: dict[tuple, SharedFc12TunableRunner] = {} + + +def release_shared_fc12_cache(device_index: int, stream_handle: int | None = None) -> None: + """Drain and release shared-FC12 runners for one executor stream. + + Dispatches already in progress finish before teardown. New dispatches for + the same stream wait, while unrelated streams and devices remain usable. + Passing no stream handle retains the process-exit fallback that releases + every runner on the device. + """ + device_index = int(device_index) + stream_handle = None if stream_handle is None else int(stream_handle) + stream_key = _begin_shared_fc12_release(device_index, stream_handle) + + def matches(key: tuple) -> bool: + return int(key[0]) == device_index and ( + stream_handle is None or int(key[6]) == stream_handle + ) + + try: + with torch.cuda.device(device_index): + if stream_handle is None: + torch.cuda.synchronize(device_index) + else: + torch.cuda.ExternalStream(stream_handle, device=device_index).synchronize() + with _SHARED_FC12_CACHE_CONDITION: + tunable_keys = [key for key in _tunable_runners if matches(key)] + for key in tunable_keys: + _tunable_runners.pop(key)._compiled_runners.clear() + kernel_keys = [key for key in SharedFc12TunableRunner.kernel_cache if matches(key)] + for key in kernel_keys: + SharedFc12TunableRunner.kernel_cache.pop(key, None) + runner_keys = [key for key in _RUNNER_CACHE if matches(key)] + for key in runner_keys: + _RUNNER_CACHE.pop(key, None) + finally: + _end_shared_fc12_release(device_index, stream_key) + + +def run_shared_fc12( + activation: torch.Tensor, + fc1_weight: torch.Tensor, + fc1_weight_sf: torch.Tensor, + fc2_weight: torch.Tensor, + fc2_weight_sf: torch.Tensor, + *, + swiglu_limit: float | None, + sm_count: int, + max_tokens: int = 8192, +) -> torch.Tensor: + """Run one complete FC12 kernel after the existing shared-expert input quantizer. + + ``activation`` is contiguous BF16 ``[T, H]``. FP8 weights have shapes + ``[1, 2*I, H]`` (gate/up interleaved in groups of 16) and ``[1, H, I]``. + Their uint8 scales are flattened R128c4 UE8M0 ``[1, padded_rows * K/32]``. + ``H`` and ``I`` must be divisible by 128; output is BF16 ``[T, H]``. + ``max_tokens`` is a fixed model capacity, not the current batch size. + """ + if activation.dtype != torch.bfloat16 or not activation.is_cuda: + raise ValueError("shared FC12 requires CUDA BF16 activations") + if activation.ndim != 2 or not activation.is_contiguous(): + raise ValueError("shared FC12 requires contiguous [T, H] activations") + if max_tokens <= 0 or activation.shape[0] > max_tokens: + raise ValueError("shared FC12 token count exceeds its configured capacity") + hidden_size = activation.shape[1] + if fc2_weight.ndim != 3 or fc2_weight.shape[:2] != (1, hidden_size): + raise ValueError("shared FC12 weight must have shape [1, H, I]") + intermediate_size = fc2_weight.shape[2] + if hidden_size % 128 or intermediate_size % 128: + raise ValueError("shared FC12 hidden and intermediate dimensions must be multiples of 128") + if fc1_weight.shape != (1, 2 * intermediate_size, hidden_size): + raise ValueError("shared FC1 weight must have shape [1, 2*I, H]") + for weight in (fc1_weight, fc2_weight): + if weight.dtype != torch.float8_e4m3fn or not weight.is_contiguous(): + raise ValueError("shared FC12 weights must be contiguous E4M3 tensors") + for scale in (fc1_weight_sf, fc2_weight_sf): + if scale.dtype != torch.uint8 or not scale.is_contiguous(): + raise ValueError("shared FC12 scales must be contiguous uint8 tensors") + if any( + t.device != activation.device + for t in (fc1_weight, fc1_weight_sf, fc2_weight, fc2_weight_sf) + ): + raise ValueError("shared FC12 inputs and weights must be on the same CUDA device") + if activation.shape[0] == 0: + return torch.empty_like(activation) + stream = torch.cuda.current_stream(activation.device) + inputs = [activation, fc1_weight, fc1_weight_sf, fc2_weight, fc2_weight_sf] + input_signature = tuple( + (str(t.dtype), (None, hidden_size) if index == 0 else tuple(t.shape), tuple(t.stride())) + for index, t in enumerate(inputs) + ) + runner_key = ( + activation.device.index, + hidden_size, + intermediate_size, + max_tokens, + sm_count, + swiglu_limit, + stream.cuda_stream, + input_signature, + ) + active_key = _begin_shared_fc12_call(activation.device.index, stream.cuda_stream) + try: + with _SHARED_FC12_CACHE_LOCK, torch.cuda.device(activation.device): + runner = _tunable_runners.get(runner_key) + if runner is None: + runner = _tunable_runners[runner_key] = SharedFc12TunableRunner(*runner_key) + tuner = AutoTuner.get() + selected_runner, tactic = tuner.choose_one( + _SHARED_FC12_OP, [runner], runner.tuning_config, inputs + ) + if tuner.is_tuning_mode: + winners = runner._cached_winner_tactics(tuner, inputs) + if winners is not None: + runner._retain_cached_winners(winners) + return selected_runner(inputs, tactic=tactic) + finally: + _end_shared_fc12_call(active_key) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 0fdb49950ec4..60d6be2ba9de 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -936,6 +936,9 @@ def _set_warmup_effects(self, in_warmup: bool) -> None: self.moe_load_balancer.set_iter_info( enable_statistic=not in_warmup, enable_update_weights=not in_warmup ) + set_warmup = getattr(self.moe_load_balancer, "set_warmup", None) + if callable(set_warmup): + set_warmup(in_warmup) @contextmanager def warmup_scope(self) -> Iterator[None]: @@ -1105,80 +1108,91 @@ def invalidate_multimodal_encoder_cache(self) -> None: if encoder_cache is not None: encoder_cache.clear() - def cleanup(self) -> None: - """Release resources owned by this model engine. - - Tears down, in order: - - 1. The optional ``ModelLoader`` (which in turn releases any - GMS client; see :meth:`ModelLoader.cleanup`). - 2. Runner resources and executor-owned CUDA Graph captures (via - :meth:`_release_cuda_graphs`). - 3. The runner, model caller, MM item scheduler, and model references. - 4. Input processors. - - Idempotency: - Subsequent calls are no-ops (guarded by ``_cleanup_done``). - The flag is set only at the end, so a partial cleanup that - raises mid-way will be retried on the next call. - - Called from: - :meth:`__del__`, and only from there. ``PyExecutor.shutdown`` - deliberately does *not* call this: it is also invoked mid-init by - ``configure_kv_cache_capacity``, which reads ``model`` right - afterwards, so clearing ``model`` here would break it. That path - uses an explicit temporary-executor shutdown and then drops its - reference instead. - """ + def prepare_cleanup(self, + *, + release_nccl_window_owners: bool = True) -> None: + """Release local graph captures before collective model teardown.""" + if self._cleanup_done: + return + self._release_cuda_graphs( + release_nccl_window_owners=release_nccl_window_owners) + + def shutdown_moe_load_balancer(self, *, local_safe: bool) -> bool: + """Run the manager's rank-aligned teardown preflight and collectives.""" + if self._cleanup_done: + return True + moe_load_balancer = getattr(self, "moe_load_balancer", None) + if moe_load_balancer is None: + return bool(local_safe) + return moe_load_balancer.shutdown( + local_safe=bool(local_safe)) is not False + + def finish_cleanup(self) -> None: + """Release rank-local model resources after collective teardown.""" if self._cleanup_done: return - # Cleanup is not truly atomic: released CUDA/GMS resources cannot be - # rolled back. Keep each handle live until its own release succeeds, - # so a failed cleanup can be retried without double-freeing resources - # that were already released. model_loader = self.model_loader if model_loader is not None: model_loader.cleanup() self.model_loader = None - # Release runner-owned graphs before dropping the runner. Keep the - # handle available if graph release fails and cleanup is retried. - # Destruction timing is not rank-symmetric; only explicit shutdown releases owners. - self._release_cuda_graphs(release_nccl_window_owners=False) - - # The runner, caller and scheduler retain the model, so - # clearing the engine's attribute alone would leave the weights - # reachable past `release_gc()` below. + self.moe_load_balancer = None + # The runner, caller and scheduler retain the model, so clearing the + # engine's attribute alone would leave the weights reachable. self._runner = None self._model_caller = None self._mm_item_scheduler = None self._lora_model = None self._lora_owner = None self.model = None - self.input_processor = None - - # Release model weights. release_gc() self._cleanup_done = True - def __del__(self) -> None: - """Best-effort cleanup during garbage collection. - - Delegates to :meth:`cleanup`. Catches ``RuntimeError`` (which a - release step such as :meth:`_release_cuda_graphs` or - ``ModelLoader.cleanup`` may raise) and ``AttributeError`` (typical - on partially-initialized engines torn down during interpreter - shutdown when module references have already been cleared); both - are logged and swallowed because destructors cannot reliably - surface exceptions. - - This is the only production caller of :meth:`cleanup` -- see the - note there on why ``PyExecutor.shutdown`` must not call it. + def cleanup(self, *, collective: bool = True) -> None: + """Release model resources without rank-asymmetric collectives. + + Explicit terminal cleanup first releases local graph captures on every + rank, then passes that local result into the load-balancer's collective + preflight, and only then drops rank-local model state. Python finalizers + perform local cleanup only; worker-owned terminal cleanup is responsible + for collective resources. """ + if self._cleanup_done: + return + + if not collective: + # Destruction timing is not rank-symmetric; only explicit shutdown + # releases NCCL window owners. + self.prepare_cleanup(release_nccl_window_owners=False) + if getattr(self, "moe_load_balancer", None) is not None: + logger.warning( + "Skipping collective MoE load-balancer shutdown from Python finalization; " + "terminal executor shutdown must release it explicitly") + self.finish_cleanup() + return + + prepare_error = None + try: + self.prepare_cleanup() + except BaseException as error: # Preserve rank alignment before surfacing it. + prepare_error = error + + manager_safe = self.shutdown_moe_load_balancer( + local_safe=prepare_error is None) + if prepare_error is not None: + raise prepare_error + if not manager_safe: + raise RuntimeError( + "MoE load-balancer collective shutdown was not safe on every rank" + ) + self.finish_cleanup() + + def __del__(self) -> None: + """Best-effort rank-local cleanup during garbage collection.""" try: - self.cleanup() + self.cleanup(collective=False) except (RuntimeError, AttributeError) as e: logger.warning("ModelEngine cleanup failed during destruction: %s", e) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 05d965bcf045..976b027d8f90 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -505,6 +505,8 @@ class PyExecutor: # The model's executor events (RuntimeHooks.executor_events); read once at # construction. _executor_events: Optional[ExecutorEvents] = None + # Shutdown publication is leader-only, while join/finalize runs on every rank. + shutdown_all_ranks = True # Minimum number of async micro batches for async PP execution. # This is a trade-off between memory usage and performance. @@ -1120,6 +1122,8 @@ def _initialize_execution(self, resources: PyExecutorResources) -> None: self.encoder_launch_executor.submit( self._warmup_encoder_cuda_graphs_enc_dec).result() + self._handoff_rebalance_warmup_owners() + # Snapshot some cumulative KV cache counters so that stats reported to # users exclude blocks reused and missed during warmup dummy requests. if hasattr(self.kv_cache_manager, 'snapshot_warmup_baseline'): @@ -1398,6 +1402,38 @@ def _kv_connector_moves_kv(self) -> bool: connector = getattr(self, "kv_connector_manager", None) return connector is not None and not connector.capacity_only + def _handoff_rebalance_warmup_owners(self) -> None: + """Transfer model-level EPLB submitters from warmup to serving.""" + managers = {} + for engine in (self.model_engine, self.draft_model_engine): + manager = getattr(engine, "moe_load_balancer", None) + handoff = getattr(manager, "handoff_submission_owner", None) + if callable(handoff): + managers[id(manager)] = manager + if not managers: + return + + # Target, draft, and encoder warmup have all completed. One device + # drain covers every manager's compute and copy streams before the + # serving worker binds submission ownership. + torch.cuda.synchronize(self.execution_stream.device) + execution_stream_handle = int(self.execution_stream.cuda_stream) + for manager in managers.values(): + manager.handoff_submission_owner(execution_stream_handle) + + def _release_quiesced_rebalance_submission_owners(self) -> None: + """Release EPLB submitters after the worker exits and device drains.""" + execution_stream_handle = int(self.execution_stream.cuda_stream) + managers = {} + for engine in (self.model_engine, self.draft_model_engine): + manager = getattr(engine, "moe_load_balancer", None) + release = getattr(manager, "release_quiesced_submission_owner", + None) + if callable(release): + managers[id(manager)] = manager + for manager in managers.values(): + manager.release_quiesced_submission_owner(execution_stream_handle) + def _maybe_init_kv_connector_manager(self) -> None: if self.kv_connector_manager is not None: # A connector that moves no KV cannot contend with the transceiver @@ -2288,7 +2324,8 @@ def _release_execution_resources(self, logger.error("Hang detected, shutting down immediately.") # Stop the relay thread and the native progress thread from # probing while the hang detector aborts MPI. A relay-thread probe - # already in flight is not interrupted. + # already in flight is not interrupted. Do not let the outer worker + # enter model collectives on a broken rank set. self._pp_mpi_progress_stop.set() self._stop_mpi_progress_thread() return False @@ -2316,10 +2353,13 @@ def _release_execution_resources(self, encoder_launch_executor.shutdown(wait=True) self.encoder_launch_executor = None - # Stop the sampler's async worker, if it was used - if (isinstance(self.sampler, AsyncWorkerMixin) - and self.sampler.async_worker_enabled()): - self.sampler.async_worker_stop() + # Stop the sampler's async worker, if it was used. A retried + # shutdown stops it only once. + if not getattr(self, "_shutdown_sampler_complete", False): + if (isinstance(self.sampler, AsyncWorkerMixin) + and self.sampler.async_worker_enabled()): + self.sampler.async_worker_stop() + self._shutdown_sampler_complete = True self.worker_started = False if self._kv_connector_failed: @@ -2334,10 +2374,9 @@ def _release_execution_resources(self, # empty_cache), the subsequent CUDA graph teardown can trigger a # device-wide cudaErrorIllegalAddress when the driver touches metadata # for the now-freed memory regions. - for engine in (self.model_engine, self.draft_model_engine): - if engine is not None: - engine.release_executor_resources( - release_model_owned_graphs=release_model_owned_graphs) + for engine in self._terminal_model_engines: + engine.release_executor_resources( + release_model_owned_graphs=release_model_owned_graphs) # Ensure graph destruction has fully completed on device before # resource managers start freeing GPU-backed workspaces. if torch.cuda.is_available(): @@ -2351,9 +2390,10 @@ def _release_execution_resources(self, if self.virtual_memory_pools is not None: keys = list(self.virtual_memory_pools.keys()) for key in keys: - del self.virtual_memory_pools[key] - if self.dwdp_manager is not None: - self.dwdp_manager.__exit__(None, None, None) + del virtual_memory_pools[key] + dwdp_manager = getattr(self, "dwdp_manager", None) + if dwdp_manager is not None: + dwdp_manager.__exit__(None, None, None) self.dwdp_manager = None return True diff --git a/tensorrt_llm/_torch/weight_sharing/source_identity.py b/tensorrt_llm/_torch/weight_sharing/source_identity.py index 4fa8305d5f92..7ede164b78c6 100644 --- a/tensorrt_llm/_torch/weight_sharing/source_identity.py +++ b/tensorrt_llm/_torch/weight_sharing/source_identity.py @@ -376,7 +376,8 @@ def _build_backend_fingerprint(model_config: "ModelConfig") -> str: Returns: A hex digest covering attention/MoE backends, allowed GEMM - backends, fusion flags, and the all-reduce strategy. + backends, fusion flags, the all-reduce strategy, and -- only when + one is configured -- the MoE load-balancer config. """ payload = { "attn_backend": getattr(model_config, "attn_backend", None), @@ -401,6 +402,9 @@ def _build_backend_fingerprint(model_config: "ModelConfig") -> str: "use_cute_dsl_bf16_bmm": getattr(model_config, "use_cute_dsl_bf16_bmm", False), "use_cute_dsl_bf16_gemm": getattr(model_config, "use_cute_dsl_bf16_gemm", False), } + load_balancer = getattr(model_config, "moe_load_balancer", None) + if load_balancer is not None: + payload["moe_load_balancer"] = load_balancer.model_dump(mode="python") return _canonical_hash(payload) @staticmethod diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 92c3cfe1bd54..6cb34992c25c 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -16,6 +16,7 @@ import enum import gc import json +import os import time import traceback import uuid @@ -36,6 +37,7 @@ nvtx_range_debug) from ..bindings import executor as tllm from ..bindings.BuildInfo import check_cuda_architecture_supported +from ..llmapi._load_balance_env import configure_moe_launch_queues from ..llmapi.llm_args import BaseLlmArgs, ExecutorMemoryType from ..llmapi.tokenizer import TokenizerBase from ..llmapi.tracer import global_tracer @@ -98,6 +100,15 @@ def __init__( tokenizer: Optional[TokenizerBase] = None, llm_args: Optional[BaseLlmArgs] = None, ) -> None: + # Direct worker construction also bypasses BaseLLM/worker_main. + if llm_args is not None and llm_args.env_overrides: + os.environ.update(llm_args.env_overrides) + if llm_args is not None and llm_args.backend == "pytorch": + queue_overrides = configure_moe_launch_queues( + getattr(llm_args, "moe_config", None), llm_args.env_overrides) + if queue_overrides is not llm_args.env_overrides: + llm_args.env_overrides = queue_overrides + postproc_config = postproc_worker_config or PostprocWorkerConfig() super().__init__( num_postprocess_workers=postproc_config.num_postprocess_workers, @@ -1066,16 +1077,50 @@ def wakeup(self, wakeup_tags: list[str]) -> None: def shutdown(self): if self.doing_shutdown: return - else: - self.doing_shutdown = True + self.doing_shutdown = True - if self.engine is not None: - can_shutdown = getattr(self.engine, "can_shutdown", None) - if can_shutdown is None: - can_shutdown = self.engine.can_enqueue_requests - if can_shutdown(): - self.engine.shutdown() - self.engine = None + try: + if self.engine is not None: + engine = self.engine + engine_shutdown_completed = bool( + getattr(self, "_engine_shutdown_completed", False)) + can_shutdown = getattr(engine, "can_shutdown", None) + if can_shutdown is None: + can_shutdown = engine.can_enqueue_requests + # PyExecutor separates leader-only shutdown publication from the + # all-rank join/finalize path. Other backends retain their gate. + # A retry must bypass that gate once shutdown already completed. + if (engine_shutdown_completed + or getattr(engine, "shutdown_all_ranks", False) + or can_shutdown()): + shutdown_error = None + local_safe = engine_shutdown_completed + if not local_safe: + try: + engine.shutdown() + self._engine_shutdown_completed = True + local_safe = True + except BaseException as error: + shutdown_error = error + terminal_cleanup = getattr(engine, "terminal_cleanup", None) + if terminal_cleanup is not None: + try: + if getattr(engine, "shutdown_all_ranks", False): + terminal_cleanup(local_safe=local_safe) + elif local_safe: + terminal_cleanup() + except BaseException as error: + if shutdown_error is None: + shutdown_error = error + if shutdown_error is not None: + raise shutdown_error + self.engine = None + self._engine_shutdown_completed = False + except BaseException: + # Keep the engine reachable so a caller can retry a failed + # collective terminal cleanup in this long-lived worker process. + self.doing_shutdown = False + raise def get_disaggregated_params(self) -> dict: if self.engine is None or self.engine.kv_cache_transceiver is None: diff --git a/tensorrt_llm/executor/ray/executor.py b/tensorrt_llm/executor/ray/executor.py index 37fa717ec3e8..d3dd2315422b 100644 --- a/tensorrt_llm/executor/ray/executor.py +++ b/tensorrt_llm/executor/ray/executor.py @@ -364,17 +364,41 @@ def shutdown(self): # Then, shutdown the workers if hasattr(self, 'workers') and self.workers is not None: - try: + shutdown_refs = getattr(self, "_pending_shutdown_refs", None) + if shutdown_refs is None: shutdown_refs = [ worker.shutdown.remote() for worker in self.workers ] - # Add timeout to prevent indefinite hanging - ray.get(shutdown_refs, timeout=30.0) - except ray.exceptions.GetTimeoutError: - logger.warning( - "Timeout waiting for workers to shutdown after 30 seconds") - except Exception as e: - logger.warning(f"Error shutting down: {e}") + self._pending_shutdown_refs = shutdown_refs + + # Converge every rank before deciding whether actor teardown is safe. + # Keep timed-out refs for the next explicit shutdown() call: issuing a + # second collective shutdown while the first is still running can + # deadlock the rank set. + shutdown_errors = [] + timed_out = False + deadline = time.monotonic() + 30.0 + for rank, shutdown_ref in enumerate(shutdown_refs): + try: + ray.get(shutdown_ref, + timeout=max(0.0, deadline - time.monotonic())) + except ray.exceptions.GetTimeoutError as error: + timed_out = True + shutdown_errors.append((rank, error)) + except Exception as error: + shutdown_errors.append((rank, error)) + + if not timed_out: + self._pending_shutdown_refs = None + if shutdown_errors: + for rank, error in shutdown_errors: + logger.warning( + f"Ray worker {rank} shutdown failed: {error}") + # Preserve the actors and placement group. RayGPUWorker resets + # its own shutdown guard on failure, so a later explicit call can + # retry after all pending refs have converged. + self._shutdown_event.clear() + raise shutdown_errors[0][1] # The engines are already stopped by the shutdown RPC above, so # kill the actor processes explicitly instead of relying on diff --git a/tensorrt_llm/executor/ray/gpu_worker.py b/tensorrt_llm/executor/ray/gpu_worker.py index 25b124fb21ac..f2033c1d939e 100644 --- a/tensorrt_llm/executor/ray/gpu_worker.py +++ b/tensorrt_llm/executor/ray/gpu_worker.py @@ -30,6 +30,7 @@ control_action_decorator) from ... import TorchLlmArgs +from ...llmapi._load_balance_env import configure_moe_launch_queues from ...llmapi.llm_args import BaseLlmArgs, ExecutorMemoryType from ...llmapi.tokenizer import TokenizerBase from ...llmapi.utils import configure_cpu_affinity @@ -56,6 +57,15 @@ def resolve_obj_by_qualname(qualname: str) -> Any: class RayWorkerWrapper: def __init__(self, worker_cls, worker_kwargs, world_size, rank): + llm_args = worker_kwargs.get("llm_args") + if llm_args is not None and llm_args.env_overrides: + os.environ.update(llm_args.env_overrides) + if llm_args is not None and llm_args.backend == "pytorch": + queue_overrides = configure_moe_launch_queues( + getattr(llm_args, "moe_config", None), llm_args.env_overrides) + if queue_overrides is not llm_args.env_overrides: + llm_args.env_overrides = queue_overrides + self.master_address = os.environ["MASTER_ADDR"] self.world_size = world_size self.rank = rank @@ -378,43 +388,89 @@ def shutdown(self): if self.doing_shutdown: return - else: - self.doing_shutdown = True - - logger.debug(f'Worker {self.rank} shutting down...') + self.doing_shutdown = True + shutdown_succeeded = False - if hasattr(self, 'shutdown_event'): - self.shutdown_event.set() - - if getattr(self, '_postproc_pool', None) is not None: - self.shutdown_postproc_workers() - - if hasattr(self, 'rpc_server') and self.rpc_server is not None: - logger.info(f"[Rank {self.global_rank}] Shutting down RPC server") + try: + logger.debug(f'Worker {self.rank} shutting down...') + + if hasattr(self, 'shutdown_event'): + self.shutdown_event.set() + + shutdown_error = None + local_safe = True + if getattr(self, '_postproc_pool', None) is not None: + try: + self.shutdown_postproc_workers() + except BaseException as error: + shutdown_error = error + local_safe = False + + if hasattr(self, 'rpc_server') and self.rpc_server is not None: + logger.info( + f"[Rank {self.global_rank}] Shutting down RPC server") + try: + self.rpc_server.shutdown() + except Exception as e: + # Suppress errors during RPC server shutdown. These can occur + # if the server is already closed or during cleanup. + logger.debug( + f"[Rank {self.global_rank}] Suppressed error during RPC server shutdown: {e}" + ) + self.rpc_server = None + + if self.engine is not None: + engine = self.engine + engine_shutdown_completed = bool( + getattr(self, "_engine_shutdown_completed", False)) + if not engine_shutdown_completed: + try: + engine.shutdown() + self._engine_shutdown_completed = True + except BaseException as error: + if shutdown_error is None: + shutdown_error = error + local_safe = False + terminal_cleanup = getattr(engine, "terminal_cleanup", None) + if terminal_cleanup is not None: + try: + if getattr(engine, "shutdown_all_ranks", False): + terminal_cleanup(local_safe=local_safe) + elif local_safe: + terminal_cleanup() + except BaseException as error: + if shutdown_error is None: + shutdown_error = error + + if (getattr(getattr(self, "llm_args", None), "backend", None) + == "pytorch" and hasattr(self, "checkpoint_loader") + and self.checkpoint_loader is not None): + try: + self.checkpoint_loader.cleanup() + except BaseException as error: + if shutdown_error is None: + shutdown_error = error + else: + self.checkpoint_loader = None + + # Check if there are any errors from the threads before shutdown. try: - self.rpc_server.shutdown() - except Exception as e: - # Suppress errors during RPC server shutdown - # These can occur if the server is already closed or during cleanup - logger.debug( - f"[Rank {self.global_rank}] Suppressed error during RPC server shutdown: {e}" - ) - self.rpc_server = None + self._handle_background_error() + except BaseException as error: + if shutdown_error is None: + shutdown_error = error - if self.engine is not None: - self.engine.shutdown() + if shutdown_error is not None: + raise shutdown_error self.engine = None - - if (self.llm_args.backend == "pytorch" - and hasattr(self, "checkpoint_loader") - and self.checkpoint_loader is not None): - self.checkpoint_loader.cleanup() - self.checkpoint_loader = None - - # Check if there are any errors from the threads before shutdown. - self._handle_background_error() - - logger.debug(f"Worker {self.rank} shutdown done.") + self._engine_shutdown_completed = False + shutdown_succeeded = True + logger.debug(f"Worker {self.rank} shutdown done.") + finally: + if not shutdown_succeeded: + # Preserve the engine and allow an explicit retry to finish any + # model-owned collective teardown that could not complete. + self.doing_shutdown = False def _configure_device(self): configure_cpu_affinity(self.device_id) diff --git a/tensorrt_llm/executor/rpc/rpc_server.py b/tensorrt_llm/executor/rpc/rpc_server.py index 07e20b5cd729..95a65407a6da 100644 --- a/tensorrt_llm/executor/rpc/rpc_server.py +++ b/tensorrt_llm/executor/rpc/rpc_server.py @@ -417,6 +417,20 @@ async def _process_requests(self) -> None: logger_debug( f"[server] RPC Server sent response for request {req}" ) + # Worker cleanup and server exit use a two-phase + # handshake. Wake RpcWorker.main_task only after the + # successful commit response is on the wire; otherwise + # main_task could cancel this coroutine before the client + # receives its acknowledgement. + if (req.method_name == "commit_shutdown" + and response.error is None): + complete_commit = getattr( + self._instance, + "_complete_shutdown_commit", + None, + ) + if callable(complete_commit): + complete_commit() # Decrement pending count if req.method_name not in ["_rpc_shutdown", "shutdown"]: diff --git a/tensorrt_llm/executor/rpc_proxy.py b/tensorrt_llm/executor/rpc_proxy.py index 2108628b7cb3..60dd334b26c1 100644 --- a/tensorrt_llm/executor/rpc_proxy.py +++ b/tensorrt_llm/executor/rpc_proxy.py @@ -297,7 +297,11 @@ def setup_engine_remote(self): def shutdown_remote(self): logger_debug("Shutting down rpc remote", color="yellow") - self.rpc_client.shutdown().remote(need_response=False) + # Cleanup must be acknowledged before main_task is allowed to close the + # server. The second RPC commits that acknowledged all-rank result; its + # response is sent before RpcWorker wakes main_task. + self.rpc_client.shutdown().remote(need_response=True) + self.rpc_client.commit_shutdown().remote(need_response=True) def abort_request(self, request_id: int) -> None: return self.rpc_client.abort_request(request_id).remote() @@ -322,13 +326,18 @@ def shutdown(self): self._shutdown_event.set() logger_debug("Shutting down GenerationExecutorRpcProxy", color="yellow") - # 1. shutdown the rpc server (PyExecutor Rank 0 + RPC server) - self.shutdown_remote() + # Keep the server, MPI workers, and retry gate live until every rank has + # acknowledged cleanup. RpcWorker resets its engine shutdown guard on a + # failed attempt, so the caller may invoke shutdown() again explicitly. + try: + self.shutdown_remote() + except BaseException: + self._shutdown_event.clear() + raise - # 2. stop the main loop, so that no new rpc requests + # Stop the response loop only after cleanup and commit are acknowledged. if self.main_loop and self.main_loop_task_obj: logger_debug("Cancelling main loop task.", color="yellow") - # The cancel() is thread-safe try: self.main_loop.call_soon_threadsafe( self.main_loop_task_obj.cancel) @@ -336,38 +345,39 @@ def shutdown(self): logger_debug(f"Error cancelling main loop task: {e}", color="yellow") - # Only join if we're not calling from the main_loop_thread itself - # (e.g., during garbage collection in that thread) if self.main_loop_thread and threading.current_thread( ) != self.main_loop_thread: self.main_loop_thread.join(timeout=2.0) if self.main_loop_thread.is_alive(): logger.warning("Main loop thread did not exit gracefully") - # 3. shutdown the mpi session, this should wait until all the PyExecutor - # processes are shutdown + shutdown_errors = [] + # Worker futures complete only after followers observe the commit and + # rank zero has joined its RPC server. Always converge them, including + # for an externally owned pool, before allowing that pool to be reused. + for future in getattr(self, "worker_futures", []): + try: + future.result() + except BaseException as error: + shutdown_errors.append(error) + if self.mpi_session is not None: if self._owns_mpi_session: logger_debug("Shutting down mpi session", color="yellow") - self.mpi_session.shutdown() - else: - # Externally owned (shared) session: leave the pool alive, but - # wait for this executor's worker tasks to finish so the next - # LLM on the pool doesn't race with PyExecutor teardown. Block - # without a timeout, mirroring the owned path above - # (mpi_session.shutdown() also waits indefinitely); a timeout - # that expires would just reintroduce the race it prevents. - for future in getattr(self, "worker_futures", []): - try: - future.result() - except Exception as e: - logger.warning( - f"RPC worker task raised during shutdown on " - f"shared MPI session: {e}") + try: + self.mpi_session.shutdown() + except BaseException as error: + shutdown_errors.append(error) logger_debug("Mpi session shutdown", color="yellow") self.mpi_session = None - self.rpc_client.close() + try: + self.rpc_client.close() + except BaseException as error: + shutdown_errors.append(error) + + if shutdown_errors: + raise shutdown_errors[0] def __enter__(self): return self diff --git a/tensorrt_llm/executor/rpc_worker.py b/tensorrt_llm/executor/rpc_worker.py index 273192943105..28c638e4f576 100644 --- a/tensorrt_llm/executor/rpc_worker.py +++ b/tensorrt_llm/executor/rpc_worker.py @@ -19,6 +19,9 @@ from .rpc import RPCServer from .rpc_worker_mixin import RpcWorkerMixin +_RPC_SHUTDOWN_RETRY = "retry" +_RPC_SHUTDOWN_COMMIT = "commit" + class RpcWorker(RpcWorkerMixin, BaseWorker): """ @@ -75,6 +78,10 @@ def __init__( and hasattr(llm_args, 'garbage_collection_gen0_threshold') else None) self.shutdown_event = Event() + self._rpc_collective_shutdown = False + self._rpc_shutdown_attempted = False + self._rpc_shutdown_ready = False + self._rpc_shutdown_committed = False self._response_queue = Queue() self.set_result_queue(self._response_queue) @@ -92,13 +99,67 @@ def setup_engine(self): super().setup_engine() def shutdown(self): + if getattr(self, "_rpc_shutdown_committed", False): + return + logger_debug(f"[worker] RpcWorker #{mpi_rank()} is shutting down", color="yellow") - self.shutdown_event.set() - super().shutdown() + if not getattr(self, "_rpc_collective_shutdown", False): + super().shutdown() + self.shutdown_event.set() + return + + # Followers wait between attempts. Wake them only when rank zero is + # explicitly retrying a failed, fully-converged shutdown attempt. + if mpi_rank() == 0 and self._rpc_shutdown_attempted: + mpi_comm().bcast(_RPC_SHUTDOWN_RETRY, root=0) + self._rpc_shutdown_attempted = True + self._rpc_shutdown_ready = False + + local_error = None + try: + super().shutdown() + except BaseException as error: + local_error = error + + # Do not acknowledge rank zero until every rank has completed cleanup. + # A peer-only failure must be reported to the client as well, otherwise + # a shared MPI pool could be reused while that peer is still tearing down. + local_failure = (None if local_error is None else + f"{type(local_error).__name__}: {local_error}") + failures = mpi_comm().allgather(local_failure) + failures = [ + f"rank {rank}: {failure}" for rank, failure in enumerate(failures) + if failure is not None + ] + if failures: + if local_error is not None: + raise local_error + raise RuntimeError("RPC worker shutdown failed: " + + "; ".join(failures)) + + self._rpc_shutdown_ready = True logger_debug(f"[worker] RpcWorker #{mpi_rank()} is shutdown", color="yellow") + def commit_shutdown(self) -> None: + """Commit an all-rank cleanup after its RPC response was acknowledged.""" + if self._rpc_shutdown_committed: + return + if mpi_rank() != 0: + raise RuntimeError("Only RPC rank zero may commit shutdown") + if not self._rpc_shutdown_ready: + raise RuntimeError( + "RPC worker cleanup has not completed on every rank") + mpi_comm().bcast(_RPC_SHUTDOWN_COMMIT, root=0) + self._rpc_shutdown_committed = True + + def _complete_shutdown_commit(self) -> None: + """Wake main_task only after RPCServer has sent the commit response.""" + if not self._rpc_shutdown_committed: + raise RuntimeError("RPC worker shutdown was not committed") + self.shutdown_event.set() + def start(self): pass @@ -126,6 +187,7 @@ def main_task( hf_model_dir=hf_model_dir, tokenizer=tokenizer, ) + worker._rpc_collective_shutdown = True if mpi_rank() != 0: # The non-leader worker will setup the engine immediately. @@ -135,6 +197,27 @@ def main_task( f"[worker] Worker {mpi_rank()} is setting up the engine", color="yellow") worker.setup_engine() + # Rank zero publishes the shutdown request. Every follower waits + # for its broadcast sentinel, then runs the same finalize path so + # model-owned EP collectives remain rank-aligned. + worker.engine.wait_shutdown() + while True: + shutdown_error = None + try: + worker.shutdown() + except BaseException as error: + shutdown_error = error + command = mpi_comm().bcast(None, root=0) + if command == _RPC_SHUTDOWN_COMMIT: + if shutdown_error is not None: + raise shutdown_error + worker._rpc_shutdown_committed = True + worker.shutdown_event.set() + break + if command != _RPC_SHUTDOWN_RETRY: + raise RuntimeError( + f"Unexpected RPC shutdown command: {command!r}" + ) from shutdown_error else: logger_debug( diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py index d017baa0cf6c..459c8867f53c 100644 --- a/tensorrt_llm/executor/worker.py +++ b/tensorrt_llm/executor/worker.py @@ -28,6 +28,7 @@ from tensorrt_llm.logger import logger from .._utils import mpi_comm, mpi_rank, print_all_stacks +from ..llmapi._load_balance_env import configure_moe_launch_queues from ..llmapi.llm_args import BaseLlmArgs from ..llmapi.mpi_session import set_mpi_session_cpp from ..llmapi.tokenizer import TokenizerBase @@ -153,73 +154,130 @@ def shutdown(self): if self.doing_shutdown: return - else: - self.doing_shutdown = True - - logger_debug(f'Worker {mpi_rank()} shutdown...\n', "yellow") + self.doing_shutdown = True + shutdown_succeeded = False - # The server references this worker; a long-lived worker process would - # otherwise keep one server, its threads and sockets per LLM it hosted. - if self.rpc_server is not None: - self.rpc_server.shutdown() - self.rpc_server = None - - if self.engine is not None: - if self.engine.can_enqueue_requests(): - if self.await_response_thread.is_alive(): - self.await_response_thread.stop() - self.await_response_thread.join() + try: + logger_debug(f'Worker {mpi_rank()} shutdown...\n', "yellow") + + shutdown_error = None + # The server references this worker; a long-lived worker process would + # otherwise keep one server, its threads and sockets per LLM it hosted. + if self.rpc_server is not None: + try: + self.rpc_server.shutdown() + self.rpc_server = None + except BaseException as error: + shutdown_error = error + if self.engine is not None: + engine = self.engine + local_safe = bool( + getattr(self, "_engine_shutdown_completed", False)) + if not local_safe: + try: + if engine.can_enqueue_requests(): + if self.await_response_thread.is_alive(): + self.await_response_thread.stop() + self.await_response_thread.join() + + engine.shutdown() + self._engine_shutdown_completed = True + local_safe = True + except BaseException as error: + shutdown_error = error + try: + terminal_cleanup = getattr(engine, "terminal_cleanup", None) + if terminal_cleanup is not None: + if getattr(engine, "shutdown_all_ranks", False): + terminal_cleanup(local_safe=local_safe) + elif local_safe: + terminal_cleanup() + except BaseException as error: + if shutdown_error is None: + shutdown_error = error + + if (self.llm_args.backend == "pytorch" + and hasattr(self, "checkpoint_loader") + and self.checkpoint_loader is not None + and shutdown_error is None): + try: + self.checkpoint_loader.cleanup() + except BaseException as error: + shutdown_error = error + else: + self.checkpoint_loader = None + + # Keep every rank aligned before entering the next collective phase. + # A rank-local engine or checkpoint failure must prevent peers from + # advancing to symmetric-buffer release and process-group teardown. + shutdown_failures = mpi_comm().allgather( + None if shutdown_error is + None else f"{type(shutdown_error).__name__}: {shutdown_error}") + shutdown_failures = [ + f"rank {rank}: {failure}" + for rank, failure in enumerate(shutdown_failures) + if failure is not None + ] + if shutdown_failures: + if shutdown_error is not None: + raise shutdown_error + raise RuntimeError("Worker shutdown failed: " + + "; ".join(shutdown_failures)) + + # MegaMoE's NVLink symmetric-memory activation workspaces are + # rendezvoused over the EP group, so they must go before the + # destroy_process_group() below. + mega_moe = sys.modules.get(_MEGA_MOE_DEEPGEMM_MODULE) + symm_release_error = None + if mega_moe is not None: + try: + mega_moe.release_symm_buffer_cache() + except BaseException as error: + symm_release_error = error + # No rank may destroy its process group while a peer still needs that + # group to retry symmetric-buffer destruction. + symm_release_failures = mpi_comm().allgather( + None if symm_release_error is None else + f"{type(symm_release_error).__name__}: {symm_release_error}") + symm_release_failures = [ + f"rank {rank}: {failure}" + for rank, failure in enumerate(symm_release_failures) + if failure is not None + ] + if symm_release_failures: + if symm_release_error is not None: + raise symm_release_error + raise RuntimeError("MegaMoE symm-buffer release failed: " + + "; ".join(symm_release_failures)) + + # Destroy torch distributed process groups so that NCCL communicators + # are torn down cleanly before MPI session shutdown and process exit. + # This is done here (not in PyExecutor.shutdown()) because the MPI + # worker owns the process group. In the Ray path the process group + # belongs to RayWorkerWrapper and must not be destroyed by the engine. + import torch.distributed + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + # Return this rank's GPU memory to the driver. Under an external MPI + # launch (mpirun/srun, e.g. CI), the worker process is long-lived and + # shared across successive LLM instances: a new GenerationExecutorWorker + # is built for each LLM, but the OS process -- and with it the CUDA + # context and PyTorch caching allocator -- persists. + gc.collect() + torch.cuda.empty_cache() - self.engine.shutdown() self.engine = None + self._engine_shutdown_completed = False + shutdown_succeeded = True + logger_debug(f"Worker {mpi_rank()} shutdown done.\n", "yellow") - if (self.llm_args.backend == "pytorch" - and hasattr(self, "checkpoint_loader") - and self.checkpoint_loader is not None): - self.checkpoint_loader.cleanup() - self.checkpoint_loader = None - - # MegaMoE's NVLink symmetric-memory activation workspaces are - # rendezvoused over the EP group, so they must go before the - # destroy_process_group() below. Never let a failure here escape: - # doing_shutdown is already set, so a retry would no-op, and skipping - # the teardown below would leave NCCL communicators alive -- turning a - # MegaMoE-only failure into a worker teardown hang. - mega_moe = sys.modules.get(_MEGA_MOE_DEEPGEMM_MODULE) - if mega_moe is not None: - try: - mega_moe.release_symm_buffer_cache() - except Exception as e: - logger.error( - f"Failed to release MegaMoE symm buffers on shutdown: {e}") - - # Destroy torch distributed process groups so that NCCL communicators - # are torn down cleanly before MPI session shutdown and process exit. - # This is done here (not in PyExecutor.shutdown()) because the MPI - # worker owns the process group. In the Ray path the process group - # belongs to RayWorkerWrapper and must not be destroyed by the engine. - import torch.distributed - if torch.distributed.is_initialized(): - torch.distributed.destroy_process_group() - - # Return this rank's GPU memory to the driver. Under an external MPI - # launch (mpirun/srun, e.g. CI), the worker process is long-lived and - # shared across successive LLM instances: a new GenerationExecutorWorker - # is built for each LLM, but the OS process -- and with it the CUDA - # context and PyTorch caching allocator -- persists. Setting - # `self.engine = None` above is not enough to free the GPU: reference - # cycles keep the model tensors alive until a later GC, and the allocator - # holds freed blocks as "reserved" instead of returning them. Without - # this, the previous model's ~weights-sized reservation carries into the - # next LLM built in this process and can OOM its load (e.g. back-to-back - # tests in one CI shard). - gc.collect() - torch.cuda.empty_cache() - - # Check if there are any errors from the threads before shutdown. - self._handle_background_error() - - logger_debug(f"Worker {mpi_rank()} shutdown done.\n", "yellow") + # Raise background-thread errors only after the commit above, which + # every rank has reached, so a retry cannot re-enter collectives. + self._handle_background_error() + finally: + if not shutdown_succeeded: + self.doing_shutdown = False def block_subordinates(self): if self.rank != 0: @@ -247,6 +305,18 @@ def worker_main( hmac_key: bytes = b"", ) -> None: + if llm_args is not None and llm_args.env_overrides: + # MPI may have cached the parent's environment before this worker was + # spawned, so replay overrides before any rebalance state is inspected. + os.environ.update(llm_args.env_overrides) + + # Configure per-layer dynamic EPLB before MPI or CUDA state is created. + if llm_args is not None and llm_args.backend == "pytorch": + queue_overrides = configure_moe_launch_queues( + getattr(llm_args, "moe_config", None), llm_args.env_overrides) + if queue_overrides is not llm_args.env_overrides: + llm_args.env_overrides = queue_overrides + def _print_stacks(): counter = 0 while True: @@ -264,14 +334,6 @@ def _print_stacks(): mpi_comm().barrier() - if llm_args is not None and llm_args.env_overrides: - # this is needed because MPI_Init seems to cache the env at import time. - # The cached env snapshot is used to spawn workers. - # Any env overrides to the main process after tensorrt_llm import - # may not get reflected in the spawned worker process, no matter how early, - # unless we update it explicitly here. - os.environ.update(llm_args.env_overrides) - if llm_args is not None and llm_args.trust_remote_code: _init_hf_modules() diff --git a/tensorrt_llm/llmapi/_load_balance_env.py b/tensorrt_llm/llmapi/_load_balance_env.py new file mode 100644 index 000000000000..5fd77b9d4b62 --- /dev/null +++ b/tensorrt_llm/llmapi/_load_balance_env.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Configure process-start CUDA settings for opt-in MoE load balancing. + +Keep this module free of GPU imports: it runs before argument validation and +worker initialization are allowed to probe CUDA. +""" + +from __future__ import annotations + +import os +import sys +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from .llm_args import MoeConfig + + +def configure_moe_launch_queues( + moe_config: MoeConfig | None, + env_overrides: dict[str, Any] | None = None, +) -> dict[str, Any] | None: + """Configure the per-layer dynamic EPLB launch queue before CUDA setup. + + ``moe_config`` must already be validated. Standard EPLB and configurations + without a load balancer leave the environment and overrides untouched. + CUDA has no supported operation to resize existing launch queues here. + The late-init guard detects PyTorch initialization; external CUDA users + must still arrange for this setting before initializing their contexts. + """ + if moe_config is not None: + moe_config.resolve_load_balancer_compatibility() + load_balancer = getattr(moe_config, "load_balancer", None) + if getattr(load_balancer, "mode", None) != "per_layer_dynamic": + return env_overrides + + name = "CUDA_SCALE_LAUNCH_QUEUES" + torch = sys.modules.get("torch") + if os.environ.get(name) != "4x" and torch is not None and torch.cuda.is_initialized(): + raise RuntimeError( + "Per-layer dynamic EPLB requires CUDA_SCALE_LAUNCH_QUEUES=4x " + "before CUDA initialization. Restart the process with that " + "environment setting or construct the LLM before using CUDA." + ) + + overrides = dict(env_overrides or {}) + overrides[name] = "4x" + os.environ[name] = "4x" + return overrides diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index ca6bec788a23..c9bb16835e05 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -55,10 +55,11 @@ from ..logger import logger from ..sampling_params import LogitsProcessor, SamplingParams from ..scheduling_params import SchedulingParams +from ._load_balance_env import configure_moe_launch_queues from .llm_args import (ENCODER_RUNNER_MANAGED_INPUTS, TORCH_LLMARGS_EXPLICIT_DOCSTRING, TORCH_LLMARGS_REMOVED_ARGS, ExecutorMemoryType, - RuntimeMemoryStatus, TorchLlmArgs, + MoeConfig, RuntimeMemoryStatus, TorchLlmArgs, validate_token_encoder_bucket_config) from .llm_utils import CachedModelLoader, KvCacheRetentionConfig, ModelLoader from .mpi_session import MpiPoolSession, external_mpi_comm_available @@ -384,7 +385,20 @@ def __init__(self, try: env_overrides = kwargs.get("env_overrides", None) + # Rebalance validation and launch-queue setup read process state. + # Apply caller overrides before either one observes that state. self._process_env_overrides(env_overrides) + if kwargs.get("backend") == "pytorch": + # This must precede TorchLlmArgs validators and GPU probes. + moe_config = kwargs.get("moe_config") + if isinstance(moe_config, dict): + moe_config = MoeConfig(**moe_config) + kwargs["moe_config"] = moe_config + queue_overrides = configure_moe_launch_queues( + moe_config, env_overrides) + if queue_overrides is not env_overrides: + kwargs["env_overrides"] = queue_overrides + env_overrides = queue_overrides backend = kwargs.get('backend', None) if backend == "pytorch": diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 5c423261ed9d..69d8dcd9c982 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -33,7 +33,8 @@ from pydantic import AliasChoices, BaseModel, ConfigDict from pydantic import Field as PydanticField from pydantic import (NonNegativeFloat, NonNegativeInt, PositiveInt, - PrivateAttr, StrictInt, field_validator, model_validator) + PrivateAttr, StrictInt, TypeAdapter, field_validator, + model_validator) from strenum import StrEnum from transformers import PreTrainedTokenizerBase @@ -1849,61 +1850,38 @@ def to_sparse_params(self, **kwargs): uses_spcompress=self.uses_spcompress) -class MoeLoadBalancerConfig(StrictBaseModel): - """Pydantic configuration model for the Mixture of Experts (MoE) load balancer. - - This model holds configuration data (`num_slots`, etc.) as well as - runtime state (`_ep_rank`, `_ep_size`) which must be set via the - `setup()` method before use. - """ +class _MoeLoadBalancerConfigBase(StrictBaseModel): + """Common slot layout and expert-parallel runtime state for EPLB.""" num_slots: Optional[int] = None - initial_global_assignments: Optional[Dict[int, List[int]]] = Field( - default=None, - repr=False # Exclude this large dict from model representation - ) - layer_updates_per_iter: int = 0 _ep_rank: Optional[int] = PrivateAttr(default=None) _ep_size: Optional[int] = PrivateAttr(default=None) - # --- Methods --- - def setup(self, ep_rank: int, ep_size: int) -> None: - """Initializes the runtime state of the configuration. - This must be called before accessing properties like `num_local_slots`. - """ + """Bind the process-local expert-parallel rank and size.""" self._ep_rank = ep_rank self._ep_size = ep_size - - # This assertion was in the original and is critical. if self.num_slots is None: raise ValueError("`num_slots` cannot be None when calling setup().") - if self.num_slots % ep_size != 0: raise ValueError( f"`num_slots` ({self.num_slots}) must be divisible by `ep_size` ({ep_size})." ) - # --- Computed Properties --- - # These properties depend on the runtime state set by setup() - @property def ep_rank(self) -> int: - """Public accessor for the private expert parallel rank.""" if self._ep_rank is None: raise AttributeError("ep_rank is not set. Call setup() first.") return self._ep_rank @property def ep_size(self) -> int: - """Public accessor for the private expert parallel size.""" if self._ep_size is None: raise AttributeError("ep_size is not set. Call setup() first.") return self._ep_size @property def num_local_slots(self) -> int: - """Calculates the number of slots local to this rank.""" if self.num_slots is None or self._ep_size is None: raise ValueError( "Cannot calculate `num_local_slots`. " @@ -1912,7 +1890,6 @@ def num_local_slots(self) -> int: @property def slot_start(self) -> int: - """Calculates the starting global slot index for this rank.""" if self._ep_rank is None: raise ValueError( "Cannot calculate `slot_start`. Call setup() first.") @@ -1920,34 +1897,116 @@ def slot_start(self) -> int: @property def slot_end(self) -> int: - """Calculates the ending global slot index (exclusive) for this rank.""" return self.slot_start + self.num_local_slots + +class MoeLoadBalancerConfig(_MoeLoadBalancerConfigBase): + """Configure the standard asynchronous EPLB implementation.""" + + mode: Literal["standard"] = "standard" + initial_global_assignments: Optional[Dict[int, List[int]]] = Field( + default=None, + repr=False, + ) + layer_updates_per_iter: int = 0 + def get_layer_initial_global_assignments( self, layer_idx: int) -> Optional[List[int]]: - """Retrieves the initial global assignments for a specific layer.""" if self.initial_global_assignments is None: return None - if layer_idx not in self.initial_global_assignments: raise KeyError( f"layer_idx {layer_idx} not found in `initial_global_assignments`." ) - assignments = self.initial_global_assignments[layer_idx] - if self.num_slots is None: raise ValueError( "`num_slots` is not set, cannot verify assignment length.") - if len(assignments) != self.num_slots: raise ValueError( f"Assignment length ({len(assignments)}) for layer {layer_idx} " f"does not match `num_slots` ({self.num_slots}).") - return assignments +class MoePerLayerDynamicLoadBalancerConfig(_MoeLoadBalancerConfigBase): + """Configure per-layer dynamic EPLB with Halo-Q scheduling.""" + + mode: Literal["per_layer_dynamic"] = "per_layer_dynamic" + num_slots: PositiveInt + auxiliary_sms: int = Field( + default=8, + ge=8, + le=128, + description="SM budget shared serially by HALO-Q and in-switch copy.", + ) + + def validate_expert_capacity(self, num_experts: int) -> None: + """Validate that every EP rank has at least one helper slot.""" + if self._ep_size is None: + raise ValueError("Call setup() before validating expert capacity.") + if num_experts <= 0: + raise ValueError("`num_experts` must be positive.") + if num_experts % self._ep_size != 0: + raise ValueError( + f"`num_experts` ({num_experts}) must be divisible by " + f"`ep_size` ({self._ep_size}) for per-layer dynamic EPLB.") + if self.num_slots <= num_experts: + raise ValueError( + f"`num_slots` ({self.num_slots}) must exceed `num_experts` " + f"({num_experts}) for per-layer dynamic EPLB.") + + +MoeLoadBalancerConfigType = Annotated[ + Union[MoeLoadBalancerConfig, MoePerLayerDynamicLoadBalancerConfig], + PydanticField(discriminator="mode"), +] +_MOE_LOAD_BALANCER_ADAPTER = TypeAdapter(MoeLoadBalancerConfigType) + + +def _normalize_load_balancer_config(value: Any) -> Any: + """Preserve legacy standard-EPLB dictionaries that predate ``mode``.""" + if isinstance(value, dict) and "mode" not in value: + return {"mode": "standard", **value} + return value + + +def _parse_load_balancer_config(value: Any) -> MoeLoadBalancerConfigType: + return _MOE_LOAD_BALANCER_ADAPTER.validate_python( + _normalize_load_balancer_config(value)) + + +def _resolve_load_balancer_config(value: Any) -> Any: + """Resolve inline or YAML-backed EPLB input to its typed config.""" + if value is None or isinstance( + value, + (MoeLoadBalancerConfig, MoePerLayerDynamicLoadBalancerConfig)): + return value + + if isinstance(value, str): + config_path = value + if not os.path.exists(config_path): + raise FileNotFoundError( + f"MoE load balancer config file not found: {config_path}") + try: + with open(config_path) as f: + value = yaml.safe_load(f) + return _parse_load_balancer_config(value) + except Exception as e: + raise ValueError( + f"Failed to load MoE load balancer config file: {config_path}" + ) from e + + if isinstance(value, dict): + try: + return _parse_load_balancer_config(value) + except Exception as e: + raise ValueError( + f"Failed to load MoE load balancer config: {value}") from e + + return value + + _MoeBackend = Literal["AUTO", "CUTLASS", "CUTEDSL", "CUTEDSL_FC12", "TRTLLM", "DEEPGEMM", "DENSEGEMM", "VANILLA", "TRITON", "MARLIN", "MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL"] @@ -1967,10 +2026,9 @@ class MoeConfig(StrictBaseModel): "If set, at most max_num_tokens tokens will be sent to torch.ops.trtllm.fused_moe at the same time. If the number of tokens exceeds max_num_tokens, the input tensors will be split into chunks and a for loop will be used." ) - load_balancer: Optional[Union[object, str]] = Field( + load_balancer: Optional[Union[MoeLoadBalancerConfigType, str]] = Field( default=None, - description="Configuration for MoE load balancing.", - json_schema_extra={"type": "Union[MoeLoadBalancerConfig, dict, str]"}) + description="Configuration for standard or per-layer dynamic EPLB.") disable_finalize_fusion: bool = Field( default=False, @@ -1984,6 +2042,23 @@ class MoeConfig(StrictBaseModel): "Use low precision combine in MoE operations (only for NVFP4 quantization). When enabled, uses lower precision for combining expert outputs to improve performance." ) + @field_validator("load_balancer", mode="before") + @classmethod + def resolve_load_balancer_config(cls, value: Any) -> Any: + return _resolve_load_balancer_config(value) + + @model_validator(mode="after") + def resolve_load_balancer_compatibility(self) -> "MoeConfig": + """Resolve the backend required by per-layer dynamic EPLB.""" + if getattr(self.load_balancer, "mode", None) != "per_layer_dynamic": + return self + if self.backend == "AUTO": + self.backend = "MEGAMOE_CUTEDSL" + elif self.backend != "MEGAMOE_CUTEDSL": + raise ValueError("per-layer dynamic EPLB requires " + "backend='MEGAMOE_CUTEDSL'") + return self + Nvfp4Backend = Literal['cutlass', 'cublaslt', 'cutedsl', 'cuda_core', 'marlin'] @@ -7842,69 +7917,40 @@ def validate_gms_locality_domain_compat(self) -> 'TorchLlmArgs': @model_validator(mode="after") def validate_gms_moe_compat(self) -> 'TorchLlmArgs': - """Reject ``LoadFormat.GMS`` combined with a MoE load balancer. - - The ``MoeLoadBalancer``'s ``register_weight_slots_after_to_cuda`` - and ``finalize_model`` run AFTER the GMS RW pool is closed and - ``finalize_write`` has committed, so any CUDA allocations they - make land in non-GMS memory and are NOT part of the committed - layout that RO peers receive. The result is "wrong inference, - no error" on RO peers (broken MoE routing state). Failing at - config-validation time is strictly better than that silent - miscompute. - - The fix for this gap (running the MoE finalize work INSIDE - ``mem_pool_scope`` and BEFORE ``finalize_write`` so MoE - allocations are part of the committed layout) is tracked as - the (MoE, GMS) follow-up; see ``model_loader.py``'s - ``TODO(GMS-MOE-LB)`` comment. + """Reject GMS when MoE load balancing owns extra device state. - Returns: - ``self`` (Pydantic ``model_validator`` contract). - - Raises: - ValueError: When ``load_format == LoadFormat.GMS`` and - ``moe_config.load_balancer`` is set. + Load-balancer state is currently allocated after GMS commits its + shared layout, so read-only peers cannot safely restore that state. """ if (self.load_format == LoadFormat.GMS and self.moe_config is not None and self.moe_config.load_balancer is not None): raise ValueError( "LoadFormat.GMS is incompatible with moe_config.load_balancer " - "in this PR. The MoE load balancer's " - "register_weight_slots_after_to_cuda and finalize_model run " - "after the GMS pool closes and finalize_write commits, so " - "their allocations land outside the committed layout. RO " - "peers would receive a broken MoE routing state. Either " - "disable moe_config.load_balancer or use LoadFormat.AUTO. " - "Tracked as the (MoE, GMS) follow-up at " - "tensorrt_llm/_torch/pyexecutor/model_loader.py " - "(see TODO(GMS-MOE-LB)).") + "because load-balancer device state is not included in the " + "committed GMS layout. Disable moe_config.load_balancer or " + "use LoadFormat.AUTO.") return self @model_validator(mode="after") def validate_load_balancer(self) -> 'TorchLlmArgs': - if isinstance(self.moe_config.load_balancer, str): - if not os.path.exists(self.moe_config.load_balancer): - raise FileNotFoundError( - f"MoE load balancer config file not found: {self.moe_config.load_balancer}" - ) - try: - with open(self.moe_config.load_balancer) as f: - moe_load_balancer_config = yaml.safe_load(f) - self.moe_config.load_balancer = MoeLoadBalancerConfig( - **moe_load_balancer_config) - except Exception as e: - raise ValueError( - f"Failed to load MoE load balancer config file: {self.moe_config.load_balancer}" - ) from e - elif isinstance(self.moe_config.load_balancer, dict): - try: - self.moe_config.load_balancer = MoeLoadBalancerConfig( - **self.moe_config.load_balancer) - except Exception as e: - raise ValueError( - f"Failed to load MoE load balancer config: {self.moe_config.load_balancer}" - ) from e + self.moe_config.load_balancer = _resolve_load_balancer_config( + self.moe_config.load_balancer) + self.moe_config.resolve_load_balancer_compatibility() + return self + + @model_validator(mode="after") + def validate_per_layer_dynamic_eplb_cuda_graph(self) -> 'TorchLlmArgs': + """Per-layer dynamic EPLB plans cannot be captured in CUDA graphs.""" + if getattr(self.moe_config.load_balancer, "mode", + None) != "per_layer_dynamic": + return self + if self.prefill_cuda_graph_backend != PrefillCudaGraphBackend.DISABLED: + raise ValueError( + "per-layer dynamic EPLB does not support prefill CUDA graphs") + if self.cuda_graph_config is not None: + logger.warning("per-layer dynamic EPLB does not support CUDA " + "graphs; setting cuda_graph_config to None.") + self.cuda_graph_config = None return self @model_validator(mode='after') diff --git a/tensorrt_llm/llmapi/mpi_session.py b/tensorrt_llm/llmapi/mpi_session.py index 190a4723a0f0..8e60054ca13d 100644 --- a/tensorrt_llm/llmapi/mpi_session.py +++ b/tensorrt_llm/llmapi/mpi_session.py @@ -36,8 +36,13 @@ # Environment that MpiPoolSession forwards from the parent to its workers. _FORWARDED_ENV_PREFIXES = ("TRTLLM", "TLLM") -_FORWARDED_ENV_NAMES = (*_UNIFIED_CACHE_ENV_VARS, "FLASHINFER_CUBIN_DIR", - "MASTER_ADDR", "MASTER_PORT") +_FORWARDED_ENV_NAMES = ( + *_UNIFIED_CACHE_ENV_VARS, + "FLASHINFER_CUBIN_DIR", + "MASTER_ADDR", + "MASTER_PORT", + # MPI may use an environment snapshot from MPI_Init. + "CUDA_SCALE_LAUNCH_QUEUES") def _is_forwarded_env(name: str) -> bool: diff --git a/tensorrt_llm/llmapi/serialization.py b/tensorrt_llm/llmapi/serialization.py index d3cc0ce5753e..517415e7a4e0 100644 --- a/tensorrt_llm/llmapi/serialization.py +++ b/tensorrt_llm/llmapi/serialization.py @@ -99,7 +99,8 @@ "tensorrt_llm.executor.worker": ["GenerationExecutorWorker", "worker_main"], "tensorrt_llm.llmapi.llm_args": [ "_ParallelConfig", "CalibConfig", "CapacitySchedulerPolicy", - "KvCacheConfig", "SchedulerConfig", "LoadFormat", "DynamicBatchConfig" + "KvCacheConfig", "SchedulerConfig", "LoadFormat", "DynamicBatchConfig", + "MoeLoadBalancerConfig", "MoePerLayerDynamicLoadBalancerConfig" ], "tensorrt_llm.llmapi.mpi_session": ["RemoteTask"], "tensorrt_llm.llmapi.tokenizer": ["TransformersTokenizer"], diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 093c1f0ee281..4e66e3e88fad 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1071,6 +1071,30 @@ "kind": "value", "path": "moe_config.disable_finalize_fusion" }, + { + "capture_policy": "int", + "kind": "value", + "path": "moe_config.load_balancer.auxiliary_sms" + }, + { + "capture_policy": "int", + "kind": "value", + "path": "moe_config.load_balancer.layer_updates_per_iter" + }, + { + "allowed_values": [ + "standard", + "per_layer_dynamic" + ], + "capture_policy": "literal", + "kind": "categorical", + "path": "moe_config.load_balancer.mode" + }, + { + "capture_policy": "int|none", + "kind": "value", + "path": "moe_config.load_balancer.num_slots" + }, { "capture_policy": "int|none", "kind": "value", diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 046479edefb5..f73117e522fa 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -48,6 +48,7 @@ l0_cpu: - unittest/_torch/kernels/kimi_k3/test_k3_mla_decode_view.py - unittest/_torch/kernels/kimi_k3/test_k3_mla_launch.py - unittest/_torch/kernels/kimi_k3/test_k3_kda_pool_layout.py + - unittest/_torch/cute_dsl_kernels/test_megamoe_shared_fc12_autotune.py - unittest/_torch/compilation/test_auto_multi_stream.py - unittest/_torch/compilation/test_phase_selective_forward.py - unittest/_torch/compilation/test_remove_copy_pass.py::test_remove_copy_preserves_minimax_producer_outputs @@ -207,6 +208,7 @@ l0_cpu: - unittest/executor/test_forward_error_survival.py - unittest/executor/test_rpc_worker_mixin.py - unittest/executor/test_sleep_collective_rpc_guards.py + - unittest/executor/test_shutdown_retry_contracts.py - unittest/executor/test_stats_serializer.py - unittest/executor/test_spec_dec_perf_metrics.py - unittest/executor/test_stop_word_side_output_trim.py @@ -277,6 +279,7 @@ l0_cpu: - unittest/llmapi/test_whisper_suppress_tokens_processor.py - unittest/llmapi/test_kv_cache_dtype_override.py - unittest/llmapi/test_llm_args.py + - unittest/llmapi/test_gms_args.py - unittest/llmapi/test_mx_args.py - unittest/llmapi/test_custom_tokenizer_aliases.py # Reads the Muse Spark Open tokenizer files from LLM_MODELS_ROOT. diff --git a/tests/unittest/_torch/cute_dsl_kernels/test_megamoe_shared_fc12_autotune.py b/tests/unittest/_torch/cute_dsl_kernels/test_megamoe_shared_fc12_autotune.py new file mode 100644 index 000000000000..c09a0e32565a --- /dev/null +++ b/tests/unittest/_torch/cute_dsl_kernels/test_megamoe_shared_fc12_autotune.py @@ -0,0 +1,458 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU contract tests; GPU numerical/concurrent admission runs separately.""" + +import __future__ + +import ast +import inspect +import os +import unittest +from contextlib import nullcontext +from copy import deepcopy +from itertools import product +from pathlib import Path +from threading import Condition, Event, RLock, Thread +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +_ROOT = Path(__file__).resolve().parents[4] +_SOURCE = _ROOT / "tensorrt_llm/_torch/cute_dsl_kernels/megamoe_shared_fc12.py" +_AUTOTUNER = _ROOT / "tensorrt_llm/_torch/autotuner.py" + +pytestmark = pytest.mark.cpu_only + + +class Tensor: + def __init__(self, shape, dtype="bf16"): + self.shape = tuple(shape) + self.dtype = dtype + self.device = SimpleNamespace(index=0) + self.is_cuda = True + self.ndim = len(shape) + + def is_contiguous(self): + return True + + def stride(self): + strides = [] + n = 1 + for extent in reversed(self.shape): + strides.append(n) + n *= extent + return tuple(reversed(strides)) + + +def inputs(tokens): + return [ + Tensor((tokens, 7168)), + Tensor((1, 6144, 7168), "e4m3"), + Tensor((1, 1376256), "uint8"), + Tensor((1, 7168, 3072), "e4m3"), + Tensor((1, 688128), "uint8"), + ] + + +class TuningConfig(SimpleNamespace): + def __init__(self, **kwargs): + super().__init__(constraint_specs=(), distributed_tuning_strategy="independent", **kwargs) + self.distributed_tuning_strategy = SimpleNamespace(value="independent") + + +class TunableRunner: + def __call__(self, inputs, **kwargs): + return self.forward(inputs, **kwargs) + + +class ProfilingCache: + def __init__(self): + self.cache = {} + + def get_cache_key(self, op, runner, shapes, config, apply_map_to_tuning_buckets=True): + shapes = list(shapes) + if apply_map_to_tuning_buckets: + shapes[0] = (runner.token_bucket(shapes[0][0]), shapes[0][1]) + return op, type(runner).__name__, str(runner.unique_id()), tuple(shapes) + + def search_cache(self, op, runners, shapes, config, apply_map_to_tuning_buckets=True): + key = self.get_cache_key(op, runners[0], shapes, config, apply_map_to_tuning_buckets) + return (True, *self.cache[key]) if key in self.cache else (False, 0, -1, float("inf")) + + +class CpuContractTests(unittest.TestCase): + def setUp(self): + self.tuner = SimpleNamespace( + is_tuning_mode=False, + skip_dynamic_tuning_buckets=False, + _active_capture=None, + profiling_cache=ProfilingCache(), + ) + properties = SimpleNamespace(multi_processor_count=80, major=10, minor=7) + self.device_syncs = [] + self.stream_syncs = [] + self.empty_cache_calls = 0 + + def external_stream(stream_handle, *, device): + return SimpleNamespace( + synchronize=lambda: self.stream_syncs.append((device, stream_handle)) + ) + + torch = SimpleNamespace( + Tensor=Tensor, + bfloat16="bf16", + float8_e4m3fn="e4m3", + uint8="uint8", + cuda=SimpleNamespace( + get_device_properties=lambda device: properties, + device=lambda device: nullcontext(), + current_stream=lambda device: SimpleNamespace(cuda_stream=17), + ExternalStream=external_stream, + synchronize=lambda device: self.device_syncs.append(device), + empty_cache=lambda: setattr(self, "empty_cache_calls", self.empty_cache_calls + 1), + ), + empty_like=lambda x: Tensor(x.shape, x.dtype), + ) + self.ns = dict( + torch=torch, + AutoTuner=SimpleNamespace(get=lambda: self.tuner), + DynamicTensorSpec=SimpleNamespace, + TunableRunner=TunableRunner, + TuningConfig=TuningConfig, + Condition=Condition, + RLock=RLock, + deepcopy=deepcopy, + product=product, + ) + module = ast.parse(_SOURCE.read_text()) + module.body = [ + node for node in module.body if not isinstance(node, (ast.Import, ast.ImportFrom)) + ] + exec( + compile(module, str(_SOURCE), "exec", flags=__future__.annotations.compiler_flag), + self.ns, + ) + self.creations = [] + self.launches = [] + + def factory(*args): + self.creations.append(args) + launches = self.launches + + class Compiled: + def __call__(self, *tensors): + launches.append((args[-1], tensors[0].shape[0])) + return tensors[0] + + return Compiled() + + self.ns["_get_runner"] = factory + + def runner(self, sms=72, capacity=8192, clamp=10.0, stream=17): + return self.ns["SharedFc12TunableRunner"]( + 0, 7168, 3072, capacity, sms, clamp, stream, ("fixed_weight_layout",) + ) + + def test_catalog_includes_both_baselines_and_resource_decisions(self): + catalog = self.ns["shared_fc12_candidate_tactics"]() + self.assertEqual(len(catalog), 48) + self.assertEqual(len(set(catalog)), 48) + self.assertIn((1, 256, 128, "atomic_counter", True), catalog) + self.assertIn((2, 128, 128, "grid_stride", True), catalog) + rejected = catalog[0] + calls = [] + + def resource_check(*args): + calls.append(args) + if args[-1] == rejected: + raise ValueError("insufficient test resource") + + self.ns["_make_kernel"] = resource_check + for sms in (72, 80): + runner = self.runner(sms) + valid = runner.get_valid_tactics([], None) + self.assertEqual(set(valid), set(catalog) - {rejected}) + self.assertEqual(runner.get_valid_tactics([], None), valid) + self.assertEqual(len(calls), 96) + with self.assertRaises(ValueError): + self.ns["shared_fc12_tactic_descriptor"]((1, 256, 512, "atomic_counter", True)) + + def test_bucket_and_persistent_key_separation(self): + r = self.runner() + self.assertEqual( + [r.token_bucket(x) for x in (1, 2, 128, 129, 8191, 8192)], + [1, 128, 128, 256, 8192, 8192], + ) + self.assertEqual(self.runner(capacity=300).buckets, (1, 128, 256, 300)) + self.assertNotEqual(r.unique_id(), self.runner(80).unique_id()) + self.assertNotEqual(r.unique_id(), self.runner(clamp=None).unique_id()) + self.assertNotEqual(r.unique_id(), self.runner(capacity=4096).unique_id()) + self.assertEqual(r.unique_id(), self.runner(stream=18).unique_id()) + for invalid in (0, 8193): + with self.assertRaises(ValueError): + r.token_bucket(invalid) + + def test_preparation_is_not_fallback_and_dynamic_t_reuses_compilation(self): + r = self.runner() + self.tuner.is_tuning_mode = True + self.assertIsNone(r.forward(inputs(8192), do_preparation=True)) + self.assertEqual(self.creations, []) + with self.assertRaises(RuntimeError): + r.forward(inputs(8192), tactic=-1) + tactic = (2, 128, 128, "atomic_counter", True) + for tokens in (8192, 129, 4096, 256, 8192): + x = inputs(tokens) + self.assertIs(r.forward(x, tactic=tactic), x[0]) + self.assertEqual(len(self.creations), 1) + self.assertEqual([tokens for _, tokens in self.launches], [8192, 129, 4096, 256, 8192]) + self.tuner.is_tuning_mode = False + r.forward(inputs(256), tactic=-1) + self.assertEqual(self.creations[-1][-1], (1, 256, 128, "atomic_counter", True)) + + def test_complete_profile_cache_prunes_only_loser_workspaces(self): + r = self.runner(capacity=300) + winners = ( + (1, 256, 128, "atomic_counter", True), + (2, 128, 128, "grid_stride", True), + ) + loser = (1, 64, 256, "grid_stride", False) + unrelated = (*r._runner_args[:-1], 18, loser) + for tactic in (*winners, loser): + cached = object() + r._compiled_runners[tactic] = cached + key = (*r._runner_args, tactic) + self.ns["_RUNNER_CACHE"][key] = cached + r.kernel_cache[key] = cached + self.ns["_RUNNER_CACHE"][unrelated] = object() + r.kernel_cache[unrelated] = object() + + profiles = [] + for index, bucket in enumerate(r.buckets): + shapes = tuple(t.shape for t in inputs(bucket)) + tactic = winners[index % len(winners)] + key = self.tuner.profiling_cache.get_cache_key( + "trtllm::megamoe_shared_fc12", + r, + shapes, + r.tuning_config, + apply_map_to_tuning_buckets=False, + ) + self.tuner.profiling_cache.cache[key] = (0, tactic, 0.1) + profiles.append(SimpleNamespace(get_opt_shapes=lambda shapes=shapes: shapes)) + self.tuner._optimization_profiles = lambda config, tensors: profiles + + retained = r._cached_winner_tactics(self.tuner, inputs(300)) + self.assertEqual(retained, set(winners)) + r._retain_cached_winners(retained) + self.assertEqual(set(r._compiled_runners), set(winners)) + for tactic in winners: + key = (*r._runner_args, tactic) + self.assertIn(key, self.ns["_RUNNER_CACHE"]) + self.assertIn(key, r.kernel_cache) + loser_key = (*r._runner_args, loser) + self.assertNotIn(loser_key, self.ns["_RUNNER_CACHE"]) + self.assertNotIn(loser_key, r.kernel_cache) + self.assertIn(unrelated, self.ns["_RUNNER_CACHE"]) + self.assertIn(unrelated, r.kernel_cache) + self.assertEqual(self.stream_syncs, [(0, 17)]) + self.assertEqual(self.empty_cache_calls, 1) + r._retain_cached_winners(retained) + self.assertEqual(self.empty_cache_calls, 1) + + def test_incomplete_profile_cache_does_not_prune(self): + r = self.runner(capacity=300) + winner = (1, 256, 128, "atomic_counter", True) + loser = (1, 64, 256, "grid_stride", False) + for tactic in (winner, loser): + cached = object() + r._compiled_runners[tactic] = cached + key = (*r._runner_args, tactic) + self.ns["_RUNNER_CACHE"][key] = cached + r.kernel_cache[key] = cached + profiles = [] + for bucket in r.buckets: + shapes = tuple(t.shape for t in inputs(bucket)) + profiles.append(SimpleNamespace(get_opt_shapes=lambda shapes=shapes: shapes)) + first_key = self.tuner.profiling_cache.get_cache_key( + "trtllm::megamoe_shared_fc12", + r, + profiles[0].get_opt_shapes(), + r.tuning_config, + apply_map_to_tuning_buckets=False, + ) + self.tuner.profiling_cache.cache[first_key] = (0, winner, 0.1) + self.tuner._optimization_profiles = lambda config, tensors: profiles + + before_compiled = dict(r._compiled_runners) + before_runners = dict(self.ns["_RUNNER_CACHE"]) + before_kernels = dict(r.kernel_cache) + self.assertIsNone(r._cached_winner_tactics(self.tuner, inputs(300))) + self.assertEqual(r._compiled_runners, before_compiled) + self.assertEqual(self.ns["_RUNNER_CACHE"], before_runners) + self.assertEqual(r.kernel_cache, before_kernels) + self.assertEqual(self.stream_syncs, []) + self.assertEqual(self.empty_cache_calls, 0) + + def test_skipped_dynamic_buckets_do_not_prune(self): + r = self.runner(capacity=300) + winner = (1, 256, 128, "atomic_counter", True) + loser = (1, 64, 256, "grid_stride", False) + for tactic in (winner, loser): + cached = object() + r._compiled_runners[tactic] = cached + key = (*r._runner_args, tactic) + self.ns["_RUNNER_CACHE"][key] = cached + r.kernel_cache[key] = cached + self.tuner.skip_dynamic_tuning_buckets = True + + self.assertIsNone(r._cached_winner_tactics(self.tuner, inputs(300))) + self.assertEqual(set(r._compiled_runners), {winner, loser}) + self.assertEqual(self.stream_syncs, []) + self.assertEqual(self.empty_cache_calls, 0) + + def test_public_dispatch_obeys_chooser_and_reuses_compilation(self): + tactic = (2, 64, 256, "grid_stride", False) + choices = [] + + def choose(op, runners, config, tensors): + runner = runners[0] + key = self.tuner.profiling_cache.get_cache_key( + op, runner, tuple(t.shape for t in tensors), config + ) + self.tuner.profiling_cache.cache[key] = (0, tactic, 0.25) + choices.append(key) + return runner, tactic + + self.tuner.choose_one = choose + outputs = [ + self.ns["run_shared_fc12"](*inputs(tokens), swiglu_limit=10.0, sm_count=72) + for tokens in (129, 200) + ] + + self.assertEqual(len(choices), 2) + self.assertEqual(choices[0], choices[1]) + self.assertEqual(len(self.creations), 1) + self.assertEqual([output.shape[0] for output in outputs], [129, 200]) + self.assertEqual([tokens for _, tokens in self.launches], [129, 200]) + + def test_terminal_release_drops_all_device_runner_references(self): + first = self.runner(stream=17) + second = self.runner(stream=18) + first._compiled_runners[(1, 256, 128, "atomic_counter", True)] = object() + second._compiled_runners[(2, 128, 128, "grid_stride", True)] = object() + first_key = (0, 7168, 3072, 8192, 72, 10.0, 17, ("layout",)) + second_key = (0, 7168, 3072, 8192, 72, 10.0, 18, ("layout",)) + self.ns["_tunable_runners"].update({first_key: first, second_key: second}) + self.ns["_RUNNER_CACHE"].update({first_key: object(), second_key: object()}) + first.kernel_cache.update({first_key: object(), second_key: object()}) + + self.ns["release_shared_fc12_cache"](0) + + self.assertEqual(self.ns["_tunable_runners"], {}) + self.assertEqual(self.ns["_RUNNER_CACHE"], {}) + self.assertEqual(first.kernel_cache, {}) + self.assertEqual(first._compiled_runners, {}) + self.assertEqual(second._compiled_runners, {}) + self.assertEqual(self.device_syncs, [0]) + self.assertEqual(self.stream_syncs, []) + + def test_terminal_release_is_scoped_to_the_executor_stream(self): + first = self.runner(stream=17) + second = self.runner(stream=18) + first._compiled_runners[(1, 256, 128, "atomic_counter", True)] = object() + second._compiled_runners[(2, 128, 128, "grid_stride", True)] = object() + first_key = (0, 7168, 3072, 8192, 72, 10.0, 17, ("layout",)) + second_key = (0, 7168, 3072, 8192, 72, 10.0, 18, ("layout",)) + self.ns["_tunable_runners"].update({first_key: first, second_key: second}) + self.ns["_RUNNER_CACHE"].update({first_key: object(), second_key: object()}) + first.kernel_cache.update({first_key: object(), second_key: object()}) + + self.ns["release_shared_fc12_cache"](0, 17) + + self.assertNotIn(first_key, self.ns["_tunable_runners"]) + self.assertIn(second_key, self.ns["_tunable_runners"]) + self.assertNotIn(first_key, self.ns["_RUNNER_CACHE"]) + self.assertIn(second_key, self.ns["_RUNNER_CACHE"]) + self.assertNotIn(first_key, first.kernel_cache) + self.assertIn(second_key, first.kernel_cache) + self.assertEqual(first._compiled_runners, {}) + self.assertNotEqual(second._compiled_runners, {}) + self.assertEqual(self.device_syncs, []) + self.assertEqual(self.stream_syncs, [(0, 17)]) + + def test_overlapping_terminal_releases_are_serialized(self): + first = self.ns["_begin_shared_fc12_release"](0, 17) + acquired = Event() + + def acquire_same_scope(): + second = self.ns["_begin_shared_fc12_release"](0, 17) + acquired.set() + self.ns["_end_shared_fc12_release"](0, second) + + thread = Thread(target=acquire_same_scope) + thread.start() + blocked = not acquired.wait(0.05) + self.ns["_end_shared_fc12_release"](0, first) + + self.assertTrue(blocked) + self.assertTrue(acquired.wait(1.0)) + thread.join(timeout=1.0) + self.assertFalse(thread.is_alive()) + + def test_original_autotuner_primes_each_cached_tactic_before_serving(self): + source = ast.parse(_AUTOTUNER.read_text()) + autotuner_class = next( + n for n in source.body if isinstance(n, ast.ClassDef) and n.name == "AutoTuner" + ) + method = deepcopy( + next( + n + for n in autotuner_class.body + if isinstance(n, ast.FunctionDef) and n.name == "_prime_cached_tactics" + ) + ) + namespace = dict( + os=os, + inspect=inspect, + nvtx_range=lambda name: nullcontext(), + DistributedTuningStrategy=SimpleNamespace(MERGE=object()), + ) + exec( + compile( + ast.Module(body=[method], type_ignores=[]), + str(_AUTOTUNER), + "exec", + flags=__future__.annotations.compiler_flag, + ), + namespace, + ) + r = self.runner() + tactics = [(2, 128, 128, "grid_stride", True), (1, 256, 128, "atomic_counter", True)] + tensors = [inputs(256), inputs(8192)] + profiles = [ + SimpleNamespace(get_opt_shapes=lambda x=x: tuple(t.shape for t in x)) for x in tensors + ] + for x, tactic in zip(tensors, tactics): + key = self.tuner.profiling_cache.get_cache_key( + "trtllm::megamoe_shared_fc12", r, tuple(t.shape for t in x), r.tuning_config + ) + self.tuner.profiling_cache.cache[key] = (0, tactic, 0.1) + self.tuner._optimization_profiles = lambda config, x: profiles + self.tuner._prepare_input_tensors = lambda profile, x: inputs( + profile.get_opt_shapes()[0][0] + ) + self.tuner._primed_cached_tactics = set() + self.tuner.is_tuning_mode = True + r.tuning_config.inputs_pre_hook = None + with patch.dict(os.environ, {"TLLM_AUTOTUNER_PRIME_CACHED_TACTICS": "1"}): + for _ in range(2): + namespace["_prime_cached_tactics"]( + self.tuner, "trtllm::megamoe_shared_fc12", [r], r.tuning_config, inputs(8192) + ) + self.assertEqual(len(self.creations), 2) + self.assertEqual([t for t, _ in self.launches], tactics) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index 4bf0aba78e03..484580955b72 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -46,6 +46,7 @@ DeepseekV4DecoderLayer, DeepseekV4ForCausalLM, DeepseekV4Gate, + DeepseekV4MoE, DeepseekV4MTP, DeepseekV4WeightLoader, _copy_deepseek_v4_fused_a_weight_scale, @@ -678,6 +679,124 @@ def test_deepseek_v4_nvfp4_mixed_precision_config(): ) +@pytest.mark.parametrize( + ("quant_config_dict", "exclude_modules", "expected"), + [ + (None, None, True), + ( + {"model.layers.4.mlp.shared_experts.gate_up_proj": QuantConfig(quant_algo=None)}, + None, + False, + ), + ( + {"model.layers.4.mlp.shared_experts.down_proj": QuantConfig(quant_algo=None)}, + None, + False, + ), + (None, ["*shared_experts.gate_up_proj"], False), + (None, ["*shared_experts.down_proj"], False), + ( + { + "model.layers.4.mlp.shared_experts.gate_proj": QuantConfig( + quant_algo=QuantAlgo.FP8_BLOCK_SCALES, group_size=128 + ), + "model.layers.4.mlp.shared_experts.up_proj": QuantConfig(quant_algo=None), + }, + None, + True, + ), + # The actual shared gate/up module is one fused Linear: its first + # gate_proj/gate_up_proj override wins and up-only overrides are ignored. + ( + { + "model.layers.4.mlp.shared_experts.gate_up_proj": QuantConfig(quant_algo=None), + "model.layers.4.mlp.shared_experts.gate_proj": QuantConfig( + quant_algo=QuantAlgo.FP8_BLOCK_SCALES, group_size=128 + ), + }, + None, + False, + ), + ( + { + "model.layers.4.mlp.shared_experts.gate_proj": QuantConfig( + quant_algo=QuantAlgo.FP8_BLOCK_SCALES, group_size=128 + ), + "model.layers.4.mlp.shared_experts.gate_up_proj": QuantConfig(quant_algo=None), + }, + None, + True, + ), + (None, ["*shared_experts.gate_proj"], False), + (None, ["*shared_experts.up_proj"], False), + ( + { + "model.layers.4.mlp.shared_experts.gate_up_proj": QuantConfig( + quant_algo=QuantAlgo.FP8_BLOCK_SCALES, group_size=64 + ), + "model.layers.4.mlp.shared_experts.down_proj": QuantConfig( + quant_algo=QuantAlgo.FP8_BLOCK_SCALES, group_size=64 + ), + }, + None, + False, + ), + ], +) +def test_deepseek_v4_shared_fc12_uses_final_projection_quantization( + quant_config_dict, exclude_modules, expected +): + model_config = ModelConfig( + pretrained_config=DeepseekV4Config(), + quant_config=QuantConfig( + quant_algo=QuantAlgo.FP8_BLOCK_SCALES, + group_size=128, + exclude_modules=exclude_modules, + ), + quant_config_dict=quant_config_dict, + ) + + gate_up_quant, down_quant = DeepseekV4MoE._get_shared_expert_projection_quant_configs( + model_config, layer_idx=4 + ) + + assert DeepseekV4MoE._shared_fc12_quantization_supported(gate_up_quant, down_quant) is expected + + +@pytest.mark.parametrize(("supported", "expected"), [(False, False), (True, True)]) +def test_deepseek_v4_shared_fc12_uses_backend_capability(supported, expected): + backend = SimpleNamespace( + capabilities=SimpleNamespace(supports_per_layer_dynamic_eplb=supported) + ) + + assert DeepseekV4MoE._shared_fc12_backend_supported(backend) is expected + + +def test_deepseek_v4_shared_fc12_rejects_backend_without_capabilities(): + assert not DeepseekV4MoE._shared_fc12_backend_supported(SimpleNamespace()) + + +def test_deepseek_v4_shared_fc12_ignores_up_only_layerwise_override(): + global_quant_config = QuantConfig(quant_algo=None) + fp8_quant_config = QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES, group_size=128) + model_config = ModelConfig( + pretrained_config=DeepseekV4Config(), + quant_config=global_quant_config, + quant_config_dict={ + "model.layers.4.mlp.shared_experts.up_proj": fp8_quant_config, + "model.layers.4.mlp.shared_experts.down_proj": fp8_quant_config, + }, + ) + + gate_up_quant, down_quant = DeepseekV4MoE._get_shared_expert_projection_quant_configs( + model_config, layer_idx=4 + ) + + assert gate_up_quant is global_quant_config + assert down_quant is fp8_quant_config + assert not DeepseekV4MoE._shared_fc12_quantization_supported(gate_up_quant, down_quant) + + def test_deepseek_v4_routed_moe_quant_config_from_mxfp4_header(tmp_path, monkeypatch): monkeypatch.setattr("tensorrt_llm._torch.model_config.get_sm_version", lambda: 100) tensor_name = "layers.0.ffn.experts.0.w1.weight" diff --git a/tests/unittest/_torch/moe/kernels/parallel/test_quantization_sm_budget.py b/tests/unittest/_torch/moe/kernels/parallel/test_quantization_sm_budget.py new file mode 100644 index 000000000000..9e77ff37b6ce --- /dev/null +++ b/tests/unittest/_torch/moe/kernels/parallel/test_quantization_sm_budget.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +from utils.util import isSM100Family + +pytestmark = pytest.mark.skipif( + not isSM100Family(), reason="Requires supported SM100-family quantization kernels" +) + + +def _assert_same_bytes(actual, expected): + for actual_tensor, expected_tensor in zip(actual, expected): + assert actual_tensor.shape == expected_tensor.shape + assert actual_tensor.dtype == expected_tensor.dtype + torch.testing.assert_close( + actual_tensor.view(torch.uint8), expected_tensor.view(torch.uint8), atol=0, rtol=0 + ) + + +def _input(rows): + generator = torch.Generator(device="cuda").manual_seed(17) + return torch.randn((rows, 7168), device="cuda", dtype=torch.bfloat16, generator=generator) + + +def _rows_for_case(case, reserved_sms): + if case == "empty": + return 0 + if case == "small": + return 129 + if case == "capacity": + return 8192 + available_sms = torch.cuda.get_device_properties("cuda").multi_processor_count - reserved_sms + grid_edge = 4 * available_sms + return grid_edge if case == "grid_edge" else grid_edge + 1 + + +@pytest.mark.parametrize("row_case", ["small", "grid_edge", "grid_edge_plus_one", "capacity"]) +@pytest.mark.parametrize("swizzled", [False, True]) +@pytest.mark.parametrize("reserved_sms", [0, 8]) +def test_fp4_quantize_sm_budget_matches_default(row_case, swizzled, reserved_sms): + x = _input(_rows_for_case(row_case, reserved_sms)) + scale = torch.ones((1,), device="cuda", dtype=torch.float32) + expected = torch.ops.trtllm.fp4_quantize(x, scale, 16, False, swizzled) + actual = torch.ops.trtllm.fp4_quantize.sm_budget( + x, scale, 16, False, swizzled, reserved_sms=reserved_sms + ) + _assert_same_bytes(actual, expected) + + +@pytest.mark.parametrize( + "row_case", ["empty", "small", "grid_edge", "grid_edge_plus_one", "capacity"] +) +@pytest.mark.parametrize("reserved_sms", [0, 8]) +def test_shared_fp8_quantize_sm_budget_matches_default(row_case, reserved_sms): + x = _input(_rows_for_case(row_case, reserved_sms)) + expected = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0(x) + actual = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0.sm_budget( + x, reserved_sms=reserved_sms + ) + _assert_same_bytes(actual, expected) + + +@pytest.mark.parametrize("op", ["fp4", "fp8"]) +@pytest.mark.parametrize("budget", ["negative", "all_sms"]) +def test_quantize_sm_budget_rejects_invalid_budget(op, budget): + x = _input(129) + reserved_sms = ( + -1 + if budget == "negative" + else torch.cuda.get_device_properties(x.device).multi_processor_count + ) + with pytest.raises(RuntimeError, match="reserved_sms"): + if op == "fp4": + scale = torch.ones((1,), device="cuda", dtype=torch.float32) + torch.ops.trtllm.fp4_quantize.sm_budget( + x, scale, 16, False, False, reserved_sms=reserved_sms + ) + else: + torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0.sm_budget(x, reserved_sms=reserved_sms) + + +def test_shared_fp8_quantize_sm_budget_rejects_legacy_layout(): + with pytest.raises(RuntimeError, match="R128c4"): + torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0.sm_budget( + _input(129), False, reserved_sms=8 + ) + + +@pytest.mark.parametrize("op", ["fp4", "fp8"]) +@pytest.mark.parametrize("reserved_sms", [0, 8]) +def test_quantize_sm_budget_fake_and_cuda_graph(op, reserved_sms): + x = _input(129) + if op == "fp4": + scale = torch.ones((1,), device="cuda", dtype=torch.float32) + quantize = torch.ops.trtllm.fp4_quantize.sm_budget + args = (x, scale, 16, False, False) + else: + quantize = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0.sm_budget + args = (x, True) + kwargs = {"reserved_sms": reserved_sms} + torch.library.opcheck(quantize, args, kwargs, test_utils=("test_schema", "test_faketensor")) + expected = quantize(*args, **kwargs) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = quantize(*args, **kwargs) + graph.replay() + _assert_same_bytes(actual, expected) diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py new file mode 100644 index 000000000000..6ecf315192ba --- /dev/null +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -0,0 +1,1765 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fast CPU contracts for the Dynamic EPLB serving path.""" + +from __future__ import annotations + +import ast +import dataclasses +import functools +import importlib.util +import math +import os +import sys +from contextlib import nullcontext +from pathlib import Path +from types import SimpleNamespace +from typing import Any, List, Optional, Tuple + +import pytest + +pytestmark = pytest.mark.cpu_only + +_ROOT = Path(__file__).resolve().parents[4] +_TORCH_ROOT = _ROOT / "tensorrt_llm/_torch" +_QUEUE = "CUDA_SCALE_LAUNCH_QUEUES" + + +def _extract_class(path: Path, name: str, methods: set[str] | None, namespace: dict): + tree = ast.parse(path.read_text()) + original = next( + node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == name + ) + if methods is None: + selected = original + else: + body = [ + node + for node in original.body + if isinstance(node, ast.FunctionDef) and node.name in methods + ] + assert {node.name for node in body} == methods + selected = ast.ClassDef(name=name, bases=[], keywords=[], body=body, decorator_list=[]) + module = ast.Module( + body=[ + ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), + selected, + ], + type_ignores=[], + ) + exec(compile(ast.fix_missing_locations(module), str(path), "exec"), namespace) + return namespace[name] + + +def test_shared_slot_provider_uses_model_owned_pool_state(): + source = _TORCH_ROOT / "moe/fused_moe/mega_moe/rebalance_live_arena.py" + provider_cls = _extract_class( + source, + "SharedSlotArenaProvider", + {"__init__"}, + {"Any": Any, "Optional": Optional, "Tuple": Tuple}, + ) + first_owner = SimpleNamespace( + _rebalance_shared_slot_pools={}, + _rebalance_shared_slot_pool_order=[], + ) + second_owner = SimpleNamespace( + _rebalance_shared_slot_pools={}, + _rebalance_shared_slot_pool_order=[], + ) + + first = provider_cls(first_owner) + same_model = provider_cls(first_owner) + second = provider_cls(second_owner) + + assert first._pools is same_model._pools + assert first._pool_order is same_model._pool_order + assert first._pools is not second._pools + with pytest.raises(TypeError, match="model-local pool state"): + provider_cls(SimpleNamespace()) + + +def test_shared_slot_provider_pool_key_includes_complete_bundle(): + source = _TORCH_ROOT / "moe/fused_moe/mega_moe/rebalance_live_arena.py" + + class FakePool: + instances = [] + + def __init__(self, **kwargs): + self.bundle = kwargs["bundle"] + self._closed = False + self.group_sizes = () + self.instances.append(self) + + def build_layer(self, comm): + return object(), {}, object(), 0 + + provider_cls = _extract_class( + source, + "SharedSlotArenaProvider", + {"__init__", "build_hierarchical_live_arena"}, + { + "Any": Any, + "Dict": dict, + "Optional": Optional, + "Tuple": Tuple, + "_SharedSlotPool": FakePool, + }, + ) + + @dataclasses.dataclass(frozen=True) + class Plane: + name: str + nbytes: int + descriptor: str + + @dataclasses.dataclass(frozen=True) + class Bundle: + hidden: int + intermediate: int + planes: tuple[Plane, ...] + + # Equal names and byte sizes but different view descriptors must not share. + first_bundle = Bundle(160, 80, (Plane("weight", 12800, "storage-a"),)) + second_bundle = Bundle(160, 80, (Plane("weight", 12800, "storage-b"),)) + owner = SimpleNamespace( + _rebalance_shared_slot_pools={}, + _rebalance_shared_slot_pool_order=[], + ) + kwargs = dict( + world=8, + rank=0, + home_count=4, + helper_count=1, + device=0, + comm=object(), + ) + first_provider = provider_cls(owner) + first_provider.build_hierarchical_live_arena(bundle=first_bundle, **kwargs) + with pytest.raises(RuntimeError, match="single-geometry"): + first_provider.build_hierarchical_live_arena(bundle=second_bundle, **kwargs) + provider_cls(owner).build_hierarchical_live_arena(bundle=second_bundle, **kwargs) + + assert len(owner._rebalance_shared_slot_pools) == 2 + assert [pool.bundle for pool in FakePool.instances] == [first_bundle, second_bundle] + + +@pytest.mark.parametrize( + ( + "hidden", + "intermediate", + "expected_total_bytes", + "expected_stride_bytes", + "expected_planes", + ), + [ + ( + 7168, + 3072, + 37158924, + 37159424, + ( + ( + "mega_fc1_weight", + 0, + 22020096, + 512, + ((3, 6144, 3584), (22020096, 3584, 1), "uint8", 1, 16), + ( + (3, 3584, 6144), + (22020096, 1, 3584), + "float4_e2m1fn_x2", + 1, + 16, + ), + ), + ( + "mega_fc1_weight_sf", + 22020096, + 2752512, + 512, + ((3, 2752512), (2752512, 1), "uint8", 1, 16), + ((3, 2752512), (2752512, 1), "float8_e4m3fn", 1, 16), + ), + ( + "mega_fc2_weight", + 24772608, + 11010048, + 512, + ((3, 7168, 1536), (11010048, 1536, 1), "uint8", 1, 16), + ( + (3, 1536, 7168), + (11010048, 1, 1536), + "float4_e2m1fn_x2", + 1, + 16, + ), + ), + ( + "mega_fc2_weight_sf", + 35782656, + 1376256, + 512, + ((3, 1376256), (1376256, 1), "uint8", 1, 16), + ((3, 1376256), (1376256, 1), "float8_e4m3fn", 1, 16), + ), + ( + "fc31_alpha", + 37158912, + 4, + 16, + ((3,), (1,), "float32", 4, 4), + ((3,), (1,), "float32", 4, 4), + ), + ( + "fc2_alpha", + 37158928, + 4, + 16, + ((3,), (1,), "float32", 4, 4), + ((3,), (1,), "float32", 4, 4), + ), + ( + "fc1_norm_const", + 37158944, + 4, + 16, + ((3,), (1,), "float32", 4, 4), + ((3,), (1,), "float32", 4, 4), + ), + ), + ), + ( + 160, + 80, + 24332, + 25088, + ( + ( + "mega_fc1_weight", + 0, + 12800, + 512, + ((3, 160, 80), (12800, 80, 1), "uint8", 1, 16), + ( + (3, 80, 160), + (12800, 1, 80), + "float4_e2m1fn_x2", + 1, + 16, + ), + ), + ( + "mega_fc1_weight_sf", + 12800, + 3072, + 512, + ((3, 3072), (3072, 1), "uint8", 1, 16), + ((3, 3072), (3072, 1), "float8_e4m3fn", 1, 16), + ), + ( + "mega_fc2_weight", + 15872, + 6400, + 512, + ((3, 160, 40), (6400, 40, 1), "uint8", 1, 16), + ( + (3, 40, 160), + (6400, 1, 40), + "float4_e2m1fn_x2", + 1, + 16, + ), + ), + ( + "mega_fc2_weight_sf", + 22528, + 2048, + 512, + ((3, 2048), (2048, 1), "uint8", 1, 16), + ((3, 2048), (2048, 1), "float8_e4m3fn", 1, 16), + ), + ( + "fc31_alpha", + 24576, + 4, + 16, + ((3,), (1,), "float32", 4, 4), + ((3,), (1,), "float32", 4, 4), + ), + ( + "fc2_alpha", + 24592, + 4, + 16, + ((3,), (1,), "float32", 4, 4), + ((3,), (1,), "float32", 4, 4), + ), + ( + "fc1_norm_const", + 24608, + 4, + 16, + ((3,), (1,), "float32", 4, 4), + ((3,), (1,), "float32", 4, 4), + ), + ), + ), + ], + ids=["production", "padding-sensitive"], +) +def test_vendored_live_plane_abi_matches_independent_goldens( + hidden, + intermediate, + expected_total_bytes, + expected_stride_bytes, + expected_planes, +): + path = _TORCH_ROOT / "cute_dsl_kernels/megamoe_scheduler/sami/geometry.py" + module_name = f"_live_plane_geometry_{hidden}_{intermediate}" + spec = importlib.util.spec_from_file_location(module_name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + try: + spec.loader.exec_module(module) + finally: + sys.modules.pop(module_name, None) + + bundle = module.BundleLayout.create(hidden=hidden, intermediate=intermediate) + observed = tuple( + ( + plane.name, + plane.offset, + plane.nbytes, + plane.alignment, + plane.storage_metadata(3), + plane.kernel_metadata(3), + ) + for plane in bundle.planes + ) + assert observed == expected_planes + assert bundle.total_bytes == expected_total_bytes + assert bundle.stride_bytes == expected_stride_bytes + + +def test_ep_collective_adapter_uses_the_resolved_process_group(): + source = _TORCH_ROOT / "moe/fused_moe/mega_moe/rebalance_live_arena.py" + tree = ast.parse(source.read_text()) + selected = next( + node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "_EpComm" + ) + + calls = [] + + class FakeProcessGroup: + def __init__(self, rank, world, global_ranks): + self.rank = rank + self.world = world + self.global_ranks = global_ranks + + class FakeDistributed: + @staticmethod + def get_rank(*, group): + calls.append(("rank", group)) + return group.rank + + @staticmethod + def get_world_size(*, group): + calls.append(("world", group)) + return group.world + + @staticmethod + def get_global_rank(group, group_rank): + calls.append(("global_rank", group, group_rank)) + return group.global_ranks[group_rank] + + @staticmethod + def barrier(*, group): + calls.append(("barrier", group)) + + @staticmethod + def broadcast_object_list(values, *, src, group): + calls.append(("bcast", group, src)) + values[0] = ("from", src) + + @staticmethod + def all_gather_object(values, value, *, group): + calls.append(("allgather", group, value)) + for rank in range(len(values)): + values[rank] = (rank, value) + + namespace = { + "Any": Any, + "_disable_current_modes": nullcontext, + "torch": SimpleNamespace(distributed=FakeDistributed), + } + module = ast.Module( + body=[ + ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), + selected, + ], + type_ignores=[], + ) + exec(compile(ast.fix_missing_locations(module), str(source), "exec"), namespace) + process_group = FakeProcessGroup(rank=1, world=3, global_ranks=(11, 13, 17)) + comm = namespace["_EpComm"](process_group) + comm.validate_geometry(expected_rank=1, expected_world=3) + assert comm.allgather("handle") == [(0, "handle"), (1, "handle"), (2, "handle")] + assert comm.bcast("payload", 2) == ("from", 17) + comm.barrier() + assert all(call[1] is process_group for call in calls) + with pytest.raises(ValueError, match="rank/size differs"): + comm.validate_geometry(expected_rank=0, expected_world=3) + with pytest.raises(ValueError, match="broadcast root"): + comm.bcast("payload", 3) + + +def test_model_engine_warmup_uses_load_balancer_public_api(): + events = [] + engine_cls = _extract_class( + _TORCH_ROOT / "pyexecutor/model_engine.py", + "PyTorchModelEngine", + {"is_warmup", "moe_load_balancer_iter_info"}, + {"set_moe_comm_warmup": lambda value: events.append(("comm", value))}, + ) + manager = SimpleNamespace( + enable_statistic=True, + enable_update_weights=True, + set_iter_info=lambda **kwargs: events.append(("iter", kwargs)), + set_warmup=lambda value: events.append(("warmup", value)), + ) + engine = engine_cls() + engine.moe_load_balancer = manager + + engine.is_warmup = True + + assert events == [ + ("comm", True), + ("iter", {"enable_statistic": False, "enable_update_weights": False}), + ("warmup", True), + ] + + +def test_model_engine_cleanup_shuts_manager_after_graphs_while_model_is_alive(): + events = [] + engine_cls = _extract_class( + _TORCH_ROOT / "pyexecutor/model_engine.py", + "PyTorchModelEngine", + { + "prepare_cleanup", + "shutdown_moe_load_balancer", + "finish_cleanup", + "cleanup", + }, + {"release_gc": lambda: events.append("release_gc")}, + ) + engine = engine_cls() + model = object() + + class Manager: + def shutdown(self, *, local_safe): + assert engine.model is model + events.append(("manager.shutdown", local_safe)) + return local_safe + + engine._cleanup_done = False + engine.model_loader = None + engine._release_cuda_graphs = lambda: events.append("release_graphs") + engine.moe_load_balancer = Manager() + engine._runner = object() + engine._model_caller = object() + engine._mm_item_scheduler = object() + engine.model = model + engine.input_processor = object() + + engine.cleanup() + engine.cleanup() + + assert events == [ + "release_graphs", + ("manager.shutdown", True), + "release_gc", + ] + assert engine.moe_load_balancer is None + assert engine.model is None + + +def test_model_engine_prepare_failure_participates_in_unsafe_preflight(): + events = [] + engine_cls = _extract_class( + _TORCH_ROOT / "pyexecutor/model_engine.py", + "PyTorchModelEngine", + { + "prepare_cleanup", + "shutdown_moe_load_balancer", + "finish_cleanup", + "cleanup", + }, + {"release_gc": lambda: events.append("release_gc")}, + ) + + class Manager: + def shutdown(self, *, local_safe): + events.append(("manager.shutdown", local_safe)) + return local_safe + + engine = engine_cls() + engine._cleanup_done = False + engine.model_loader = None + engine._release_cuda_graphs = lambda: (_ for _ in ()).throw( + RuntimeError("graph cleanup failed") + ) + engine.moe_load_balancer = Manager() + engine._runner = object() + engine._model_caller = object() + engine._mm_item_scheduler = object() + engine.model = object() + engine.input_processor = object() + + with pytest.raises(RuntimeError, match="graph cleanup failed"): + engine.cleanup() + + assert events == [("manager.shutdown", False)] + assert engine.moe_load_balancer is not None + assert engine.model is not None + + +def test_model_engine_finalizer_never_enters_collective_cleanup(): + calls = [] + engine_cls = _extract_class( + _TORCH_ROOT / "pyexecutor/model_engine.py", + "PyTorchModelEngine", + {"__del__"}, + {"logger": SimpleNamespace(warning=lambda *args: calls.append(("warning", args)))}, + ) + engine = engine_cls() + engine.cleanup = lambda *, collective: calls.append(("cleanup", collective)) + + engine.__del__() + + assert calls == [("cleanup", False)] + + +def test_executor_terminal_cleanup_is_three_phase_and_stream_scoped(monkeypatch): + events = [] + module_name = "tensorrt_llm._torch.cute_dsl_kernels.megamoe_shared_fc12" + monkeypatch.setitem( + sys.modules, + module_name, + SimpleNamespace( + release_shared_fc12_cache=lambda device, stream_handle=None: events.append( + ("cache", device, stream_handle) + ) + ), + ) + executor_cls = _extract_class( + _TORCH_ROOT / "pyexecutor/py_executor.py", + "PyExecutor", + {"terminal_cleanup"}, + { + "__name__": "tensorrt_llm._torch.pyexecutor.py_executor", + "__package__": "tensorrt_llm._torch.pyexecutor", + }, + ) + + class Engine: + def __init__(self, name): + self.name = name + + def prepare_cleanup(self): + events.append(("prepare", self.name)) + + def shutdown_moe_load_balancer(self, *, local_safe): + events.append(("collective", self.name, local_safe)) + return local_safe + + def finish_cleanup(self): + events.append(("finish", self.name)) + + executor = executor_cls() + executor.execution_stream = SimpleNamespace(device=SimpleNamespace(index=3), cuda_stream=99) + executor._terminal_model_engines = (Engine("first"), Engine("second")) + + executor.terminal_cleanup() + executor.terminal_cleanup() + + assert events == [ + ("prepare", "first"), + ("prepare", "second"), + ("collective", "first", True), + ("collective", "second", True), + ("finish", "first"), + ("finish", "second"), + ("cache", 3, 99), + ] + assert executor._terminal_model_engines == () + + +def test_executor_terminal_cleanup_propagates_prepare_failure_to_all_managers(monkeypatch): + events = [] + executor_cls = _extract_class( + _TORCH_ROOT / "pyexecutor/py_executor.py", + "PyExecutor", + {"terminal_cleanup"}, + { + "__name__": "tensorrt_llm._torch.pyexecutor.py_executor", + "__package__": "tensorrt_llm._torch.pyexecutor", + }, + ) + + class Engine: + def __init__(self, name, fail=False): + self.name = name + self.fail = fail + + def prepare_cleanup(self): + events.append(("prepare", self.name)) + if self.fail: + raise RuntimeError(f"{self.name} graph failure") + + def shutdown_moe_load_balancer(self, *, local_safe): + events.append(("collective", self.name, local_safe)) + return local_safe + + def finish_cleanup(self): + events.append(("finish", self.name)) + + executor = executor_cls() + executor.execution_stream = SimpleNamespace(device=SimpleNamespace(index=3), cuda_stream=99) + executor._terminal_model_engines = (Engine("first", fail=True), Engine("second")) + + with pytest.raises(RuntimeError, match="first graph failure"): + executor.terminal_cleanup() + + assert events == [ + ("prepare", "first"), + ("prepare", "second"), + ("collective", "first", False), + ("collective", "second", False), + ] + assert len(executor._terminal_model_engines) == 2 + + +def test_pyexecutor_shutdown_tail_failure_is_retryable(): + events = [] + + class AsyncWorkerMixin: + pass + + class RetryingPool(dict): + def __init__(self): + super().__init__(pool=object()) + self.delete_attempts = 0 + + def __delitem__(self, key): + self.delete_attempts += 1 + events.append(("pool_delete", self.delete_attempts)) + if self.delete_attempts == 1: + raise RuntimeError("pool delete failed") + super().__delitem__(key) + + class Sampler(AsyncWorkerMixin): + def __init__(self): + self.stop_attempts = 0 + + def async_worker_enabled(self): + return True + + def async_worker_stop(self): + self.stop_attempts += 1 + events.append(("sampler_stop", self.stop_attempts)) + if self.stop_attempts == 1: + raise RuntimeError("sampler stop failed") + + class Engine: + def _release_cuda_graphs(self): + events.append("release_graphs") + + class DwdpManager: + def __init__(self): + self.exit_attempts = 0 + + def __exit__(self, *args): + self.exit_attempts += 1 + events.append(("dwdp_exit", self.exit_attempts)) + if self.exit_attempts == 1: + raise RuntimeError("dwdp exit failed") + + executor_cls = _extract_class( + _TORCH_ROOT / "pyexecutor/py_executor.py", + "PyExecutor", + {"shutdown", "_finish_shutdown"}, + { + "AsyncWorkerMixin": AsyncWorkerMixin, + "torch": SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: False)), + }, + ) + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/base_worker.py", + "BaseWorker", + {"shutdown"}, + {}, + ) + + engine = Engine() + executor = executor_cls() + executor.model_engine = engine + executor.draft_model_engine = None + executor.shutdown_all_ranks = True + executor.can_shutdown = lambda: True + executor.executor_request_queue = SimpleNamespace( + enqueue_shutdown_request=lambda: events.append("enqueue") + ) + executor.shutdown_event = SimpleNamespace(wait=lambda: events.append("wait")) + executor._profile_manager = SimpleNamespace(cleanup=lambda: events.append("profile_cleanup")) + executor.hang_detector = SimpleNamespace(detected=lambda: False) + executor.worker_thread = SimpleNamespace(join=lambda: events.append("worker_join")) + executor.kv_connector_manager = None + executor.dist = SimpleNamespace(pp_size=1) + executor._shutdown_sleep_wakeup_listeners = lambda: events.append("listener_shutdown") + executor.encoder_launch_executor = None + executor.worker_started = True + executor.resource_manager = SimpleNamespace( + resource_managers={ + "manager": SimpleNamespace(shutdown=lambda: events.append("resource_shutdown")) + } + ) + executor.virtual_memory_pools = RetryingPool() + executor.sampler = Sampler() + executor.dwdp_manager = DwdpManager() + + def terminal_cleanup(*, local_safe): + events.append(("terminal_cleanup", local_safe)) + if local_safe: + assert executor.model_engine is None + assert executor._terminal_model_engines == (engine,) + executor._terminal_model_engines = () + + executor.terminal_cleanup = terminal_cleanup + worker = worker_cls() + worker.doing_shutdown = False + worker.engine = executor + + for message in ("pool delete failed", "sampler stop failed", "dwdp exit failed"): + with pytest.raises(RuntimeError, match=message): + worker.shutdown() + assert executor.model_engine is engine + assert executor._terminal_model_engines == (engine,) + assert executor._shutdown_runtime_complete + assert worker.engine is executor + assert not worker.doing_shutdown + + worker.shutdown() + + assert events == [ + "enqueue", + "wait", + "profile_cleanup", + "worker_join", + "listener_shutdown", + "release_graphs", + "resource_shutdown", + ("pool_delete", 1), + ("terminal_cleanup", False), + ("pool_delete", 2), + ("sampler_stop", 1), + ("terminal_cleanup", False), + ("sampler_stop", 2), + ("dwdp_exit", 1), + ("terminal_cleanup", False), + ("dwdp_exit", 2), + ("terminal_cleanup", True), + ] + assert executor.virtual_memory_pools == {} + assert executor.model_engine is None + assert executor.draft_model_engine is None + assert executor.dwdp_manager is None + assert executor._shutdown_sampler_complete + assert executor._terminal_model_engines == () + assert worker.engine is None + + +class _Routes: + def __init__( + self, + rows, + trace, + *, + topk=6, + dtype="int32", + device="cuda:0", + contiguous=True, + ): + self.shape = (rows, topk) + self.ndim = 2 + self.dtype = dtype + self.device = device + self._contiguous = contiguous + self.trace = trace + + def is_contiguous(self): + return self._contiguous + + def record_stream(self, stream): + self.trace.append(("lifetime", stream.cuda_stream)) + + def __getitem__(self, rows): + return _Routes(rows.stop, self.trace) + + +class _StageTensor: + def __init__(self, name, trace, shape=(8, 6)): + self.name = name + self.trace = trace + self.shape = shape + + def __getitem__(self, index): + return _StageTensor(self.name, self.trace, self.shape) + + def view(self, dtype): + return self + + def copy_(self, other, **kwargs): + self.trace.append("copy:" + self.name) + + def zero_(self): + self.trace.append("zero:" + self.name) + + def fill_(self, value): + self.trace.append("fill:" + self.name) + + +def test_launch_queue_is_enabled_only_for_per_layer_dynamic_mode(monkeypatch): + path = _ROOT / "tensorrt_llm/llmapi/_load_balance_env.py" + spec = importlib.util.spec_from_file_location("load_balance_env_under_test", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + initialized = False + fake_torch = SimpleNamespace( + cuda=SimpleNamespace(is_initialized=lambda: initialized), + ) + process_queue = os.environ.get(_QUEUE) + test_environ = os.environ.copy() + test_environ.pop(_QUEUE, None) + monkeypatch.setattr(module, "os", SimpleNamespace(environ=test_environ)) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + + def config(mode="per_layer_dynamic"): + return SimpleNamespace( + load_balancer=SimpleNamespace(mode=mode), + resolve_load_balancer_compatibility=lambda: None, + ) + + original = {"KEEP": "value"} + configured = module.configure_moe_launch_queues(config(), original) + assert configured == {"KEEP": "value", _QUEUE: "4x"} + assert configured is not original + assert original == {"KEEP": "value"} + assert test_environ[_QUEUE] == "4x" + assert module.configure_moe_launch_queues(config(), configured) == configured + + for off in (None, config(mode="standard")): + test_environ.pop(_QUEUE, None) + before = {"KEEP": "value"} + assert module.configure_moe_launch_queues(off, before) is before + assert _QUEUE not in test_environ + + initialized = True + before = {"KEEP": "value"} + assert module.configure_moe_launch_queues(config(mode="standard"), before) is before + assert _QUEUE not in test_environ + + with pytest.raises(RuntimeError, match="before CUDA initialization"): + module.configure_moe_launch_queues(config(), original) + assert original == {"KEEP": "value"} + assert _QUEUE not in test_environ + assert os.environ.get(_QUEUE) == process_queue + + +def test_launch_queue_setup_precedes_runtime_initialization(): + def call_name(node): + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + return f"{call_name(node.value)}.{node.attr}" + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "super" + ): + return "super()" + return "" + + def calls(relative, function_name, class_name=None): + tree = ast.parse((_ROOT / relative).read_text()) + owner = tree + if class_name is not None: + owner = next( + node + for node in tree.body + if isinstance(node, ast.ClassDef) and node.name == class_name + ) + function = next( + node + for node in owner.body + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + and node.name == function_name + ) + result = {} + for node in ast.walk(function): + if isinstance(node, ast.Call): + name = call_name(node.func) + result[name] = min(result.get(name, node.lineno), node.lineno) + return result + + base = calls("tensorrt_llm/llmapi/llm.py", "__init__", "BaseLLM") + assert ( + base["self._process_env_overrides"] + < base["configure_moe_launch_queues"] + < base["llm_args_cls"] + ) + assert base["configure_moe_launch_queues"] < base["get_device_count"] + + worker_main = calls("tensorrt_llm/executor/worker.py", "worker_main") + assert ( + worker_main["os.environ.update"] + < worker_main["configure_moe_launch_queues"] + < worker_main["worker_cls"] + ) + assert worker_main["configure_moe_launch_queues"] < worker_main["mpi_comm"] + + worker = calls("tensorrt_llm/executor/base_worker.py", "__init__", "BaseWorker") + assert ( + worker["os.environ.update"] + < worker["configure_moe_launch_queues"] + < worker["super().__init__"] + ) + + ray = calls("tensorrt_llm/executor/ray/gpu_worker.py", "__init__", "RayWorkerWrapper") + assert ( + ray["os.environ.update"] + < ray["configure_moe_launch_queues"] + < ray["torch.cuda.device_count"] + < ray["torch.cuda.set_device"] + ) + + +def test_routed_megamoe_direct_inputs_and_full_grid_contract(): + source = _TORCH_ROOT / "moe/fused_moe/mega_moe/mega_moe_cute_dsl.py" + tree = ast.parse(source.read_text()) + launch = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "cute_dsl_megamoe_nvfp4_blackwell" + ) + arguments = {keyword.arg: keyword.value for keyword in launch.keywords} + + for name, caster in (("activation", "_as_nvfp4"), ("activation_sf", "_as_fp8_sf")): + value = arguments[name] + assert isinstance(value, ast.Call) + assert isinstance(value.func, ast.Name) and value.func.id == caster + + for name in ("fc1_weight", "fc1_weight_sf", "fc2_weight", "fc2_weight_sf"): + value = arguments[name] + assert isinstance(value, ast.Call) + assert isinstance(value.func, ast.Attribute) + assert value.func.attr == "_memo_dtype_view" + + assert "reserved_sms" not in arguments + + +def test_gpu_direct_tma_fails_closed_on_stale_scheduler_plan(): + source = ( + _ROOT / "cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu" + ).read_text() + begin = source.index("__global__ void tma_copy_gpu_direct_kernel") + end = source.index("struct MegamoeTmaCopyState", begin) + kernel = source[begin:end] + + status_check = kernel.index("statusOk &= direct.status[field] == 0") + epoch_check = kernel.index("schedulerPlanReady = statusOk && planEpoch == schedulerEpoch") + rejection = kernel.index("if (!schedulerPlanReady)") + trap = kernel.index('asm volatile("trap;"', rejection) + publish = kernel.index("publish_plan_channel_warp") + build = kernel.index("megamoe_tma_build_gpu_plan") + ready = kernel.index("finish_cta_and_notify") + + assert status_check < epoch_check < rejection <= trap < publish < build < ready + + submit = source[source.index("megamoe_tma_copy_submit_gpu_direct") :] + assert "tma_copy_gpu_direct_kernel<<<" in submit + assert "cudaLaunchAttributeProgrammaticStreamSerialization" not in submit + assert "cudaLaunchKernelEx" not in submit + + +def test_halo_q_rendezvous_resets_arrivals_and_wraps_epoch_safely(): + source = ( + _ROOT / "cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.cu" + ).read_text() + begin = source.index("__device__ bool grid_rendezvous") + end = source.index("__device__ void local_histogram_pass", begin) + rendezvous = source[begin:end] + + assert "using Epoch = std::uint32_t;" in source + assert "return observed != expected && expected - observed < (Epoch{1} << 31);" in source + arrival = rendezvous.index("fetch_add_acq_rel_device(sync, 1)") + last_cta = rendezvous.index("old == static_cast(blocks - 1)") + reset = rendezvous.index("store_relaxed_device(sync, 0)") + publish = rendezvous.index("store_release_device(sync + 1, target)") + wait = rendezvous.index("while (epoch_before(observed, target))") + assert arrival < last_cta < reset < publish < wait + + mask = (1 << 32) - 1 + + def epoch_before(observed: int, expected: int) -> bool: + delta = (expected - observed) & mask + return observed != expected and delta < (1 << 31) + + epochs = (0xFFFFFFFE, 0xFFFFFFFF, 0x00000000, 0x00000001) + arrivals = 0 + for observed, expected in zip(epochs, epochs[1:]): + for _ in range(128): + old = arrivals + arrivals += 1 + if old == 127: + arrivals = 0 + assert arrivals == 0 + assert epoch_before(observed, expected) + assert not epoch_before(expected, observed) + assert not epoch_before(expected, expected) + + +def test_direct_submit_orders_generations_and_defers_the_route_wait(): + trace = [] + state = SimpleNamespace(device=0, thread=7) + main = SimpleNamespace(cuda_stream=44, priority=0) + copy = SimpleNamespace(cuda_stream=19, priority=-1) + cuda = SimpleNamespace( + current_device=lambda: state.device, + current_stream=lambda _: main, + nvtx=SimpleNamespace(range=lambda _: nullcontext()), + ) + torch = SimpleNamespace(cuda=cuda, Tensor=_Routes, int32="int32") + source = _TORCH_ROOT / "moe/fused_moe/mega_moe/rebalance_slot_scheduler.py" + namespace = {"torch": torch, "threading": SimpleNamespace(get_ident=lambda: state.thread)} + group_cls = _extract_class(source, "RebalanceSlotSchedulerGroup", None, namespace) + lease_cls = _extract_class(source, "_LiveBankLeaseProvider", None, namespace) + group = group_cls.__new__(group_cls) + for name in ( + "_generation", + "_finished_generation", + "_scheduled_generation", + "_route_wait_generation", + "plan_calls", + ): + setattr(group, name, 0) + group._plan_part = None + group._pending_release = None + group._closed = False + group._owner_thread_id = None + group._execution_stream_handle = None + group.device = 0 + group.topk = 6 + group.max_tokens_per_rank = 8192 + group.copy_stream = group.scheduler_stream = copy + group._stream_handle = 19 + group._input_ready_handle = 1 + group._plan_ready_handle = 2 + group._consumer_done_handle = 3 + group._consumer_done = SimpleNamespace(cuda_event=3) + + def driver_call(*entry): + trace.append(entry) + return (0,) + + group._driver = SimpleNamespace( + CUresult=SimpleNamespace(CUDA_SUCCESS=0), + cuEventRecord=lambda event, stream: driver_call("record", event, stream), + cuStreamWaitEvent=lambda stream, event, flags: driver_call("wait", stream, event), + ) + outputs = SimpleNamespace(physical_slot_ids=_Routes(8192, trace)) + pending = False + + def submit(routes, handle): + nonlocal pending + assert not pending and handle == 19 + pending = True + trace.append(("HALO", routes.shape[0], id(routes))) + return outputs + + def copy_submit(result): + nonlocal pending + assert pending and result is outputs + pending = False + generation = group._scheduled_generation + trace.append(("TMA", generation)) + return SimpleNamespace(generation=generation) + + group.scheduler = SimpleNamespace(device="cuda:0", submit=submit) + group.broadcaster = SimpleNamespace( + submit=copy_submit, + release_generation_after=lambda generation, event: trace.append( + ("consumer_release", generation, event.cuda_event) + ), + mark_collective_reuse_safe=lambda provider, generation: trace.append(("lease", generation)), + ) + group.lease = lease_cls(group) + + for generation, rows in enumerate((16, 0), 1): + trace.clear() + routes = _Routes(rows, trace) + part = group.plan_schedule(routes) + physical, actual_generation = group.plan_finish(part, defer_wait=True) + assert physical.shape == (rows, 6) + assert actual_generation == generation + assert not any(entry[0] == "wait" and entry[1:] == (44, 2) for entry in trace) + group.wait_for_routes() + group.wait_for_routes() + group.finish() + assert sum(entry == ("wait", 44, 2) for entry in trace) == 1 + assert trace.index(("HALO", rows, id(routes))) < trace.index(("TMA", generation)) + assert trace.index(("TMA", generation)) < trace.index(("wait", 44, 2)) + assert trace[-2:] == [("record", 3, 44), ("consumer_release", generation, 3)] + + trace.clear() + with pytest.raises(ValueError, match="contiguous CUDA int32"): + group.plan_schedule(_Routes(8, trace, dtype="int64")) + assert trace == [] + + group._driver.cuStreamWaitEvent = lambda *args: (17,) + with pytest.raises(RuntimeError, match="cuStreamWaitEvent"): + group.plan_schedule(_Routes(8, trace)) + assert not any(entry[0] == "HALO" for entry in trace) + assert group._plan_part is None + + +def test_route_wait_occurs_after_independent_staging(): + source = _TORCH_ROOT / "moe/fused_moe/mega_moe/mega_moe_cute_dsl.py" + impl = _extract_class( + source, + "TrtllmCutedslMegaMoeNvfp4Impl", + {"_stage_inputs"}, + {"torch": SimpleNamespace(uint8="uint8")}, + ) + + def run(stage_routes): + trace = [] + owner = impl() + owner._last_staged_T = {8: 8} + owner._wait_rebalance_routes = lambda replica_plan=None: trace.append("wait") + bufs = SimpleNamespace( + **{ + key: _StageTensor(key, trace) + for key in ("topk_idx_local", "activation", "activation_sf", "topk_weights") + } + ) + owner._stage_inputs( + bufs=bufs, + x=_StageTensor("x", trace), + x_sf=_StageTensor("sf", trace), + topk_idx=_StageTensor("routes", trace), + topk_weights=_StageTensor("weights", trace), + num_tokens=4, + top_k=6, + stage_activation=True, + stage_routes=stage_routes, + ) + return trace, owner + + trace, owner = run(True) + assert trace.index("copy:activation") < trace.index("wait") + assert trace.index("copy:topk_weights") < trace.index("wait") + assert trace.index("wait") < trace.index("copy:topk_idx_local") + assert owner._last_staged_T[8] == 4 + + trace, owner = run(False) + assert "wait" not in trace + assert not any("topk_idx_local" in item for item in trace) + assert owner._last_staged_T[8] == 8 + + +def test_warmup_owner_handoff_uses_deduplicated_manager_public_api(): + events = [] + main = SimpleNamespace(cuda_stream=11, priority=0, device=0) + torch = SimpleNamespace( + cuda=SimpleNamespace( + synchronize=lambda device: events.append(("synchronize", device)), + ) + ) + executor_cls = _extract_class( + _TORCH_ROOT / "pyexecutor/py_executor.py", + "PyExecutor", + {"_handoff_rebalance_warmup_owners"}, + {"torch": torch}, + ) + + class Manager: + def __init__(self, name): + self.name = name + + def handoff_submission_owner(self, stream_handle): + events.append(("handoff", self.name, stream_handle)) + + target = Manager("target") + draft = Manager("draft") + executor = executor_cls() + executor.execution_stream = main + executor.model_engine = SimpleNamespace(moe_load_balancer=target) + executor.draft_model_engine = SimpleNamespace(moe_load_balancer=draft) + + executor._handoff_rebalance_warmup_owners() + assert events == [ + ("synchronize", 0), + ("handoff", "target", 11), + ("handoff", "draft", 11), + ] + + events.clear() + executor.draft_model_engine = SimpleNamespace(moe_load_balancer=target) + executor._handoff_rebalance_warmup_owners() + assert events == [("synchronize", 0), ("handoff", "target", 11)] + + events.clear() + executor.model_engine = SimpleNamespace(moe_load_balancer=None) + executor.draft_model_engine = None + executor._handoff_rebalance_warmup_owners() + assert events == [] + + +def test_quiesced_owner_release_uses_deduplicated_manager_public_api(): + events = [] + main = SimpleNamespace(cuda_stream=11, priority=0, device=0) + executor_cls = _extract_class( + _TORCH_ROOT / "pyexecutor/py_executor.py", + "PyExecutor", + {"_release_quiesced_rebalance_submission_owners"}, + {}, + ) + + class Manager: + def __init__(self, name): + self.name = name + + def release_quiesced_submission_owner(self, stream_handle): + events.append(("release", self.name, stream_handle)) + + target = Manager("target") + draft = Manager("draft") + executor = executor_cls() + executor.execution_stream = main + executor.model_engine = SimpleNamespace(moe_load_balancer=target) + executor.draft_model_engine = SimpleNamespace(moe_load_balancer=draft) + + executor._release_quiesced_rebalance_submission_owners() + assert events == [ + ("release", "target", 11), + ("release", "draft", 11), + ] + + events.clear() + executor.draft_model_engine = SimpleNamespace(moe_load_balancer=target) + executor._release_quiesced_rebalance_submission_owners() + assert events == [("release", "target", 11)] + + +def _autotune_helpers(torch): + source = _TORCH_ROOT / "moe/custom_ops/cute_dsl_megamoe_custom_op.py" + functions = { + "_autotune_pl_alpha", + "synthesize_profiling_topk", + "_token_back_ready_granularity_choices", + "_megamoe_tuning_op_name", + "_unpack_tactic", + "_is_pow2_in_range", + "validate_megamoe_tactic", + "_epi_flag_batch_for_tokens", + "_expand_megamoe_ready_tactics", + "enumerate_megamoe_candidate_tactics", + } + constants = { + "_NVFP4_BLOCK_SIZE", + "_SUPPORTED_MMA_TILE_M", + "_SUPPORTED_MMA_TILE_N", + "_AUTOTUNE_PL_SEED", + "_DEFAULT_TOKEN_BACK_READY_GRANULARITY", + "_TOKEN_BACK_READY_GRANULARITIES", + "_TACTIC_LEN", + "_TACTIC_LEN_V4", + "_LEGACY_TACTIC_LEN", + "_FLAG_BATCH_MAX", + "_GEOMETRIES", + "_TOKEN_BACK_MODES", + "_TOKEN_BACK_STORE_BINDING", + "_WORK_ID_MODE_CANDIDATES", + "_GROUP_HINTS", + "_FLAG_BATCHES", + "_EPI_FLAG_BATCH_SMALL", + "_EPI_FLAG_BATCH_LARGE", + "_EPI_FLAG_BATCH_TOKEN_THRESHOLD", + } + tree = ast.parse(source.read_text()) + nodes = [ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0)] + for node in tree.body: + if isinstance(node, ast.FunctionDef) and node.name in functions: + nodes.append(node) + elif isinstance(node, (ast.Assign, ast.AnnAssign)): + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + if any(isinstance(target, ast.Name) and target.id in constants for target in targets): + nodes.append(node) + namespace = { + "torch": torch, + "math": math, + "functools": functools, + "Tuple": Tuple, + "List": List, + "Optional": Optional, + "Any": Any, + "logger": SimpleNamespace(debug=lambda *args: None), + } + exec( + compile( + ast.fix_missing_locations(ast.Module(body=nodes, type_ignores=[])), str(source), "exec" + ), + namespace, + ) + return namespace + + +def test_megamoe_runner_cache_reuses_only_production_runners(): + source = _TORCH_ROOT / "moe/custom_ops/cute_dsl_megamoe_custom_op.py" + tree = ast.parse(source.read_text()) + factory_node = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.FunctionDef) and node.name == "_megamoe_get_runner" + ) + + created = [] + + class FakeRunner: + def __init__(self, **kwargs): + self.kwargs = kwargs + self._profiling_scratch = object() + self._profiling_scratch_factory = object() + created.append(self) + + namespace = { + "Sm100MegaMoENvfp4Runner": FakeRunner, + "_MEGAMOE_RUNNER_CACHE": {}, + } + module = ast.Module( + body=[ + ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), + factory_node, + ], + type_ignores=[], + ) + exec(compile(ast.fix_missing_locations(module), str(source), "exec"), namespace) + get_runner = namespace["_megamoe_get_runner"] + kwargs = dict( + world_size=8, + local_rank=2, + num_topk=6, + num_experts_per_rank=32, + hidden_size=7168, + intermediate_size_per_partition=2048, + expand_intermediate_size_per_partition=4096, + max_tokens_per_rank=8192, + output_dtype="bf16", + apply_topk_in_fc1=True, + swiglu_alpha=None, + swiglu_beta=None, + gate_up_clamp=None, + situ_beta=None, + situ_linear_beta=None, + in_kernel_fc2_reduce=False, + combine_format="bf16", + helper_expert_count=4, + ) + + production = get_runner(**kwargs, tactic_autotune=False) + production._profiling_scratch = "stale" + production._profiling_scratch_factory = "stale" + assert get_runner(**kwargs, tactic_autotune=False) is production + assert production._profiling_scratch is None + assert production._profiling_scratch_factory is None + + first_tuning = get_runner(**kwargs, tactic_autotune=True) + second_tuning = get_runner(**kwargs, tactic_autotune=True) + assert first_tuning is not second_tuning + assert len(created) == 3 + + +def test_on_autotune_uses_distinct_tactics(): + helpers = _autotune_helpers(None) + + off_key = helpers["_megamoe_tuning_op_name"](0) + on_key = helpers["_megamoe_tuning_op_name"](4) + assert off_key == "trtllm::cute_dsl_megamoe_nvfp4_blackwell" + assert on_key != off_key and "alpha=0.8" in on_key + + off = helpers["enumerate_megamoe_candidate_tactics"](32, 107) + on = helpers["enumerate_megamoe_candidate_tactics"](32, 107, load_balance=True) + assert all(helpers["_unpack_tactic"](tactic)[10] == "expert" for tactic in off) + assert {helpers["_unpack_tactic"](tactic)[10] for tactic in on} == { + "expert", + "token_tile", + } + assert any(tactic[2] is not None for tactic in on) + + +def test_on_autotune_builds_balanced_powerlaw_input(request): + torch = pytest.importorskip("torch") + original_num_threads = torch.get_num_threads() + request.addfinalizer(lambda: torch.set_num_threads(original_num_threads)) + torch.set_num_threads(4) + helpers = _autotune_helpers(torch) + kwargs = dict( + num_tokens=8192, + num_topk=6, + num_experts_per_rank=52, + world_size=8, + alpha=0.8, + device=torch.device("cpu"), + source_rank=0, + ) + routes = helpers["synthesize_profiling_topk"](**kwargs) + assert torch.equal(routes, helpers["synthesize_profiling_topk"](**kwargs)) + counts = torch.bincount(routes.reshape(-1), minlength=8 * 52).view(8, 52) + assert bool((counts[:, 1:] >= counts[:, :-1]).all()) + per_rank = counts.sum(1) + assert int(per_rank.max() - per_rank.min()) <= 1 + assert bool((counts[:, -1] > 0).all()) + + uniform = helpers["synthesize_profiling_topk"](**{**kwargs, "alpha": 0.0}) + expected = (torch.arange(8192 * 6) % (8 * 52)).view(8192, 6) + assert torch.equal(uniform, expected) + + +def test_per_layer_dynamic_shutdown_reports_in_iter_through_collective_preflight(): + manager_cls = _extract_class( + _TORCH_ROOT / "moe/fused_moe/per_layer_dynamic_eplb.py", + "PerLayerDynamicMoeLoadBalancer", + {"shutdown"}, + {}, + ) + calls = [] + + class Comm: + def allgather(self, state): + calls.append(("allgather", state)) + return [state, {"local_safe": True, "pool_count": 1}] + + class Pool: + def close_collectively(self, *, local_safe): + calls.append(("pool", local_safe)) + return True + + manager = manager_cls() + manager.is_shutdown = False + manager.in_iter = True + manager._ep_comm = Comm() + manager.ep_process_group = object() + manager._rebalance_shared_slot_pool_order = [Pool()] + + assert manager.shutdown(local_safe=True) is False + assert calls == [ + ("allgather", {"local_safe": False, "pool_count": 1}), + ] + + +def test_base_worker_finalizes_rank_follower_pyexecutor(): + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/base_worker.py", + "BaseWorker", + {"shutdown"}, + {}, + ) + events = [] + engine = SimpleNamespace( + shutdown_all_ranks=True, + can_shutdown=lambda: False, + shutdown=lambda: events.append("shutdown"), + terminal_cleanup=lambda *, local_safe: events.append(("terminal_cleanup", local_safe)), + ) + worker = worker_cls() + worker.doing_shutdown = False + worker.engine = engine + + worker.shutdown() + + assert events == ["shutdown", ("terminal_cleanup", True)] + assert worker.engine is None + + events.clear() + cleanup_attempts = 0 + gate_attempts = 0 + + def can_shutdown(): + nonlocal gate_attempts + gate_attempts += 1 + return gate_attempts == 1 + + def terminal_cleanup(): + nonlocal cleanup_attempts + cleanup_attempts += 1 + events.append(("terminal_cleanup", True)) + if cleanup_attempts == 1: + raise RuntimeError("terminal cleanup failed") + + engine = SimpleNamespace( + shutdown_all_ranks=False, + can_shutdown=can_shutdown, + shutdown=lambda: events.append("shutdown"), + terminal_cleanup=terminal_cleanup, + ) + worker = worker_cls() + worker.doing_shutdown = False + worker.engine = engine + + with pytest.raises(RuntimeError, match="terminal cleanup failed"): + worker.shutdown() + + assert events == ["shutdown", ("terminal_cleanup", True)] + assert worker.engine is engine + assert not worker.doing_shutdown + + worker.shutdown() + + assert events == [ + "shutdown", + ("terminal_cleanup", True), + ("terminal_cleanup", True), + ] + assert worker.engine is None + assert gate_attempts == 1 + + +def test_mpi_worker_terminal_cleanup_failure_is_retryable(monkeypatch): + import types + + events = [] + fake_distributed = types.ModuleType("torch.distributed") + fake_distributed.is_initialized = lambda: True + fake_distributed.destroy_process_group = lambda: events.append("destroy_process_group") + fake_torch = types.ModuleType("torch") + fake_torch.distributed = fake_distributed + fake_torch.cuda = SimpleNamespace(empty_cache=lambda: events.append("empty_cache")) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + monkeypatch.setitem(sys.modules, "torch.distributed", fake_distributed) + + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/worker.py", + "GenerationExecutorWorker", + {"shutdown"}, + { + "logger_debug": lambda *args, **kwargs: None, + "mpi_rank": lambda: 0, + "mpi_comm": lambda: SimpleNamespace(allgather=lambda value: [value]), + "sys": sys, + "_MEGA_MOE_DEEPGEMM_MODULE": "missing_test_megamoe_module", + "gc": SimpleNamespace(collect=lambda: events.append("gc")), + }, + ) + cleanup_attempts = 0 + + def terminal_cleanup(*, local_safe): + nonlocal cleanup_attempts + cleanup_attempts += 1 + events.append(("terminal_cleanup", local_safe)) + if cleanup_attempts == 1: + raise RuntimeError("terminal cleanup failed") + + engine = SimpleNamespace( + shutdown_all_ranks=True, + can_enqueue_requests=lambda: False, + shutdown=lambda: events.append("shutdown"), + terminal_cleanup=terminal_cleanup, + ) + worker = worker_cls() + worker.doing_shutdown = False + worker.engine = engine + worker.llm_args = SimpleNamespace(backend="pytorch") + worker._handle_background_error = lambda: events.append("background") + + with pytest.raises(RuntimeError, match="terminal cleanup failed"): + worker.shutdown() + + assert events == ["shutdown", ("terminal_cleanup", True)] + assert worker.engine is engine + assert not worker.doing_shutdown + + worker.shutdown() + + assert events == [ + "shutdown", + ("terminal_cleanup", True), + ("terminal_cleanup", True), + "destroy_process_group", + "gc", + "empty_cache", + "background", + ] + assert worker.engine is None + + +def test_ray_worker_postproc_failure_still_runs_unsafe_terminal_cleanup(): + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/ray/gpu_worker.py", + "RayGPUWorker", + {"shutdown"}, + { + "logger": SimpleNamespace( + debug=lambda *args, **kwargs: None, info=lambda *args, **kwargs: None + ), + }, + ) + events = [] + + def fail_postproc(): + events.append("postproc") + raise RuntimeError("postproc cleanup failed") + + worker = worker_cls() + worker.doing_shutdown = False + worker.rank = 0 + worker._postproc_pool = object() + worker.shutdown_postproc_workers = fail_postproc + engine = SimpleNamespace( + shutdown_all_ranks=True, + shutdown=lambda: events.append("shutdown"), + terminal_cleanup=lambda *, local_safe: events.append(("terminal_cleanup", local_safe)), + ) + worker.engine = engine + worker.llm_args = SimpleNamespace(backend="pytorch") + worker._handle_background_error = lambda: events.append("background") + + with pytest.raises(RuntimeError, match="postproc cleanup failed"): + worker.shutdown() + + assert events == [ + "postproc", + "shutdown", + ("terminal_cleanup", False), + "background", + ] + assert worker.engine is engine + assert not worker.doing_shutdown + + worker._postproc_pool = None + worker.shutdown() + + assert events == [ + "postproc", + "shutdown", + ("terminal_cleanup", False), + "background", + ("terminal_cleanup", True), + "background", + ] + assert worker.engine is None + + +def test_ray_worker_shutdown_preserves_first_error_and_attempts_all_cleanup(): + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/ray/gpu_worker.py", + "RayGPUWorker", + {"shutdown"}, + { + "logger": SimpleNamespace( + debug=lambda *args, **kwargs: None, info=lambda *args, **kwargs: None + ), + }, + ) + events = [] + + def fail(name): + def raise_error(): + events.append(name) + raise RuntimeError(f"{name} failed") + + return raise_error + + worker = worker_cls() + worker.doing_shutdown = False + worker.rank = 0 + worker._postproc_pool = object() + worker.shutdown_postproc_workers = fail("postproc") + worker.engine = SimpleNamespace( + shutdown_all_ranks=True, + shutdown=lambda: events.append("shutdown"), + terminal_cleanup=lambda *, local_safe: events.append(("terminal_cleanup", local_safe)), + ) + worker.llm_args = SimpleNamespace(backend="pytorch") + worker.checkpoint_loader = SimpleNamespace(cleanup=fail("checkpoint")) + worker._handle_background_error = fail("background") + + with pytest.raises(RuntimeError, match="postproc failed"): + worker.shutdown() + + assert events == [ + "postproc", + "shutdown", + ("terminal_cleanup", False), + "checkpoint", + "background", + ] + assert worker.engine is not None + assert not worker.doing_shutdown + assert worker.checkpoint_loader is not None + + +def test_ray_worker_partial_init_without_llm_args_still_checks_background(): + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/ray/gpu_worker.py", + "RayGPUWorker", + {"shutdown"}, + { + "logger": SimpleNamespace( + debug=lambda *args, **kwargs: None, info=lambda *args, **kwargs: None + ), + }, + ) + events = [] + worker = worker_cls() + worker.doing_shutdown = False + worker.rank = 0 + worker.engine = None + worker._handle_background_error = lambda: events.append("background") + + worker.shutdown() + + assert events == ["background"] + + +def test_standard_load_balancer_unsafe_preflight_skips_destructive_cleanup(): + manager_cls = _extract_class( + _TORCH_ROOT / "moe/fused_moe/moe_load_balancer.py", + "MoeLoadBalancer", + {"shutdown"}, + {}, + ) + calls = [] + + class Comm: + def allgather(self, local_safe): + calls.append(("allgather", local_safe)) + return [local_safe, True] + + def barrier(self): + calls.append("barrier") + + manager = manager_cls() + manager.is_shutdown = False + manager.in_iter = True + manager.shared_mpi_comm = Comm() + manager.single_layer_load_balancers = [] + manager.load_balancer_impl = SimpleNamespace(shutdown=lambda: calls.append("shutdown")) + + assert manager.shutdown(local_safe=True) is False + assert calls == [("allgather", False)] diff --git a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py index 7702eaaad408..e46650574760 100644 --- a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py +++ b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py @@ -2544,6 +2544,25 @@ def _run(lo: int, hi: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, to ) whole = _run(0, num_slots) + bundle = method._live_weight_bundle_layout(hidden, intermediate, expand_intermediate) + planes = {plane.name: plane for plane in bundle.planes} + for tensor, name in zip( + whole, + ( + "mega_fc1_weight", + "mega_fc1_weight_sf", + "mega_fc2_weight", + "mega_fc2_weight_sf", + ), + ): + shape, stride, dtype_name, element_size, pointer_alignment = planes[name].storage_metadata( + num_slots + ) + assert tuple(tensor.shape) == shape + assert tuple(tensor.stride()) == stride + assert tensor.dtype == getattr(torch, dtype_name) + assert tensor.element_size() == element_size + assert tensor.data_ptr() % pointer_alignment == 0 for chunk in (1, 3, 16): # 16 > num_slots: the single-call degenerate case pieces = [_run(lo, min(lo + chunk, num_slots)) for lo in range(0, num_slots, chunk)] diff --git a/tests/unittest/_torch/moe/test_megamoe_streaming_load.py b/tests/unittest/_torch/moe/test_megamoe_streaming_load.py index 2c6bb606de21..776af12b1ab1 100644 --- a/tests/unittest/_torch/moe/test_megamoe_streaming_load.py +++ b/tests/unittest/_torch/moe/test_megamoe_streaming_load.py @@ -51,7 +51,7 @@ def _load_classes() -> tuple[type[MoEWeightLoadingMode], type[NVFP4MegaMoECuteDs class _StreamingMoEModule(nn.Module): """Minimal single-rank module for the quant-method load path.""" - def __init__(self, weight_loading_mode: MoEWeightLoadingMode) -> None: + def __init__(self, weight_loading_mode: MoEWeightLoadingMode, helper_slots: int = 0) -> None: from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig super().__init__() @@ -61,6 +61,7 @@ def __init__(self, weight_loading_mode: MoEWeightLoadingMode) -> None: self.intermediate_size_per_partition = INTERMEDIATE_SIZE self.expand_intermediate_size_per_partition = 2 * INTERMEDIATE_SIZE self.expert_size_per_partition = NUM_EXPERTS + self.helper_slots = int(helper_slots) self.initial_local_expert_ids = list(range(NUM_EXPERTS)) self.tp_size = 1 self.tp_rank = 0 @@ -72,11 +73,38 @@ def __init__(self, weight_loading_mode: MoEWeightLoadingMode) -> None: # No EPLB in this test: need_load_shared_weights() must be False. self.layer_load_balancer = None + @property + def compute_slot_count(self) -> int: + return int(self.expert_size_per_partition) + self.helper_slots + def _add_raw_shared_weights_for_unmap(self, weight_tensors: list[torch.Tensor]) -> None: # Only forwards to the dynamic load balancer in production; no-op here. del weight_tensors +class _SentinelAlphaMoEModule(_StreamingMoEModule): + def __init__(self, weight_loading_mode: MoEWeightLoadingMode, helper_slots: int) -> None: + super().__init__(weight_loading_mode, helper_slots) + self.alpha_sentinels: dict[str, torch.Tensor] = {} + + def register_parameter(self, name, parameter): + if ( + name in ("fc31_alpha", "fc2_alpha") + and parameter is not None + and tuple(parameter.shape) == (NUM_EXPERTS,) + and name not in self.alpha_sentinels + ): + base = 11 if name == "fc31_alpha" else 21 + sentinel = torch.arange( + NUM_EXPERTS, + device=parameter.device, + dtype=parameter.dtype, + ).add_(base) + parameter.data.copy_(sentinel) + self.alpha_sentinels[name] = sentinel.clone() + super().register_parameter(name, parameter) + + def _w13_input_scale(expert_id: int) -> float: return 0.5 + 0.125 * expert_id @@ -178,6 +206,87 @@ def _expected_fc1_norm_const() -> torch.Tensor: ) +def test_helper_slots_widen_only_seven_kernel_planes() -> None: + mode_cls, method_cls = _load_classes() + helper_slots = 2 + module = _SentinelAlphaMoEModule(mode_cls.VANILLA, helper_slots=helper_slots) + method = method_cls() + with torch.device("cuda"): + method.create_weights(module) + + assert module.expert_size_per_partition == NUM_EXPERTS + for name in _STREAMED_PARAMS: + assert module.rebuild_tensor_metadata[name]["meta"].shape[0] == NUM_EXPERTS + + compute_slots = NUM_EXPERTS + helper_slots + kernel_planes = ( + "mega_fc1_weight", + "mega_fc1_weight_sf", + "mega_fc2_weight", + "mega_fc2_weight_sf", + "fc31_alpha", + "fc2_alpha", + "fc1_norm_const", + ) + assert all(getattr(module, name).shape[0] == compute_slots for name in kernel_planes) + + bundle = method.live_weight_plane_spec(module) + for plane in bundle.planes: + parameter = getattr(module, plane.name) + shape, stride, dtype_name, element_size, pointer_alignment = plane.storage_metadata( + compute_slots + ) + assert tuple(parameter.shape) == shape + assert tuple(parameter.stride()) == stride + assert parameter.dtype == getattr(torch, dtype_name) + assert parameter.element_size() == element_size + assert parameter.data_ptr() % pointer_alignment == 0 + + kernel_shape, kernel_stride, kernel_dtype, kernel_element_size, _ = plane.kernel_metadata( + compute_slots + ) + kernel_view = parameter.view(getattr(torch, kernel_dtype)) + if plane.name in ("mega_fc1_weight", "mega_fc2_weight"): + kernel_view = kernel_view.transpose(1, 2) + assert tuple(kernel_view.shape) == kernel_shape + assert tuple(kernel_view.stride()) == kernel_stride + assert kernel_view.dtype == getattr(torch, kernel_dtype) + assert kernel_view.element_size() == kernel_element_size + + for name, sentinel in module.alpha_sentinels.items(): + actual = getattr(module, name) + assert torch.equal(actual[:NUM_EXPERTS], sentinel) + assert torch.equal( + actual[NUM_EXPERTS:], + torch.ones(helper_slots, dtype=actual.dtype, device=actual.device), + ) + assert set(module.alpha_sentinels) == {"fc31_alpha", "fc2_alpha"} + assert module.quant_scales.fc1_global is module.fc31_alpha + assert module.quant_scales.fc2_global is module.fc2_alpha + + +@pytest.mark.parametrize("helper_slots", [0, 2]) +def test_create_weights_supports_meta_init(helper_slots: int) -> None: + """The model loader builds MoE layers under MetaInitMode. + + Any op that touches a meta tensor raises there and makes the loader rebuild + the whole model on the host, so the large derived planes must stay meta. + """ + from tensorrt_llm._torch.models.modeling_utils import MetaInitMode + + mode_cls, method_cls = _load_classes() + module = _StreamingMoEModule(mode_cls.VANILLA, helper_slots=helper_slots) + with MetaInitMode(): + method_cls().create_weights(module) + + for name in ("mega_fc1_weight", "mega_fc1_weight_sf", "mega_fc2_weight", "mega_fc2_weight_sf"): + assert getattr(module, name).is_meta + norm = module.fc1_norm_const + assert not norm.is_meta + assert norm.shape[0] == NUM_EXPERTS + helper_slots + assert torch.equal(norm, torch.ones_like(norm)) + + def test_initial_streaming_load_layer_atomic() -> None: method, module, weights = _fresh() diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index 3285b5caa1d7..e86408bdde48 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -17,6 +17,7 @@ import collections import dataclasses import importlib +import inspect import itertools import logging import os @@ -107,6 +108,7 @@ MoEProblem, MoERejection, MoERejectReason, + MoEReplicaPlan, MoEResolutionReport, MoERunContext, MoEStaticCapability, @@ -123,6 +125,10 @@ MegaMoEDeepGemm, TrtllmCutedslMegaMoeNvfp4Impl, ) +from tensorrt_llm._torch.moe.fused_moe.mega_moe.rebalance_slot_scheduler import ( + RebalanceSlotSchedulerGroup, + _LiveBankLeaseProvider, +) from tensorrt_llm._torch.moe.fused_moe.moe_resolution import ( IMPL_PRIORITY, _reject_unsupported_activation, @@ -2040,20 +2046,220 @@ def test_nvfp4_trtllm_gen_tp_accepts_logical_group_aligned_shard() -> None: assert backend.intermediate_size_per_partition % backend.quant_method.weight_alignment != 0 +@pytest.mark.cpu_only +def test_megamoe_quant_method_defines_canonical_live_plane_spec(): + method = NVFP4MegaMoECuteDslMethod() + module = SimpleNamespace( + hidden_size=7168, + intermediate_size_per_partition=3072, + expand_intermediate_size_per_partition=6144, + ) + + spec = method.live_weight_plane_spec(module) + + assert (spec.hidden, spec.intermediate, spec.expand_intermediate) == (7168, 3072, 6144) + assert tuple(plane.name for plane in spec.planes) == ( + "mega_fc1_weight", + "mega_fc1_weight_sf", + "mega_fc2_weight", + "mega_fc2_weight_sf", + "fc31_alpha", + "fc2_alpha", + "fc1_norm_const", + ) + + # Exercise the storage-to-kernel dtype/view bridge on a small legal + # geometry; production metadata is covered by the independent goldens. + bridge = method._live_weight_bundle_layout(160, 80, 160) + for plane in bridge.planes: + shape, stride, dtype_name, element_size, pointer_alignment = plane.storage_metadata(3) + tensor = torch.empty(shape, dtype=getattr(torch, dtype_name)) + assert tuple(tensor.stride()) == stride + assert tensor.element_size() == element_size + assert tensor.data_ptr() % pointer_alignment == 0 + + kernel_shape, kernel_stride, _, kernel_element_size, _ = plane.kernel_metadata(3) + kernel_view = tensor + if plane.name in ("mega_fc1_weight", "mega_fc2_weight"): + kernel_view = kernel_view.transpose(1, 2) + assert tuple(kernel_view.shape) == kernel_shape + assert tuple(kernel_view.stride()) == kernel_stride + assert kernel_view.element_size() == kernel_element_size + + +@pytest.mark.cpu_only +def test_megamoe_live_plane_spec_rejects_expand_before_allocation(): + method = NVFP4MegaMoECuteDslMethod() + module = SimpleNamespace( + hidden_size=160, + intermediate_size_per_partition=80, + expand_intermediate_size_per_partition=176, + ) + + with pytest.raises(ValueError, match="expand_intermediate == 2"): + method.live_weight_plane_spec(module) + + +@pytest.mark.cpu_only +def test_megamoe_rejects_expand_before_provider_allocation(monkeypatch): + moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) + torch.nn.Module.__init__(moe) + moe._weights_created = False + moe.expert_size_per_partition = 4 + moe._rebalance_slots_active = 1 + moe._rebalance_home_experts = None + moe.layer_idx = 0 + moe.hidden_size = 160 + moe.intermediate_size_per_partition = 80 + # The base class derives this from the expand ratio; pin an inconsistent value. + monkeypatch.setattr(MegaMoECuteDsl, "expand_intermediate_size_per_partition", 176) + moe.quant_method = NVFP4MegaMoECuteDslMethod() + moe.ep_size = 2 + allocate_provider = MagicMock() + moe._alloc_symm_provider = allocate_provider + create_arena_provider = MagicMock() + moe.layer_load_balancer = SimpleNamespace( + manager=SimpleNamespace(create_arena_provider=create_arena_provider) + ) + + with pytest.raises(ValueError, match="expand_intermediate == 2"): + moe.create_weights() + + allocate_provider.assert_not_called() + create_arena_provider.assert_not_called() + + +@pytest.mark.cpu_only +def test_megamoe_live_plane_identity_guard_detects_parameter_replacement(): + class Tensor: + def __init__(self, pointer): + self.shape = (2, 3) + self.dtype = "u8" + self.device = "cuda:0" + self.is_meta = False + self._pointer = pointer + + def stride(self): + return (3, 1) + + def data_ptr(self): + return self._pointer + + arena_view = Tensor(17) + alias = Tensor(17) + plane = SimpleNamespace(name="weight") + bundle = SimpleNamespace(planes=(plane,)) + arena = SimpleNamespace( + bundle=bundle, + local_plane_views=(arena_view,), + tekit_alias_views=(alias,), + ) + module = SimpleNamespace(_rebalance_arena=arena, weight=Tensor(17)) + method = NVFP4MegaMoECuteDslMethod() + method.live_weight_plane_spec = MagicMock(return_value=bundle) + + method.assert_live_weight_aliases(module, arena) + + module.weight = Tensor(23) + with pytest.raises(RuntimeError, match="changed storage"): + method.assert_live_weight_aliases(module, arena) + + +@pytest.mark.cpu_only +def test_megamoe_compute_slots_do_not_widen_resident_expert_axis(): + moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) + torch.nn.Module.__init__(moe) + moe.expert_size_per_partition = 8 + moe._rebalance_home_experts = 8 + moe._rebalance_slots_active = 3 + + assert moe.compute_slot_count == 11 + assert moe.expert_size_per_partition == 8 + + # Catch the old representation (expert_size_per_partition == H + S) + # instead of silently adding S a second time. + moe.expert_size_per_partition = 11 + with pytest.raises(RuntimeError, match="resident expert geometry changed"): + _ = moe.compute_slot_count + + +@pytest.mark.cpu_only +def test_megamoe_replica_plan_supplies_launch_facts_and_route_wait(): + moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) + torch.nn.Module.__init__(moe) + moe.expert_size_per_partition = 8 + moe._rebalance_home_experts = 8 + moe._rebalance_slots_active = 3 + moe._quantize_reserved_sms = 8 + wait_for_routes = MagicMock() + ready_flags = torch.empty(8, dtype=torch.int64) + plan = MoEReplicaPlan( + resident_experts_per_rank=8, + compute_slots_per_rank=11, + ready_flags=ready_flags, + ready_generation=7, + reserved_sms=8, + _wait_for_routes=wait_for_routes, + ) + + compute_slots, helper_slots, actual_flags, generation, reserved_sms = ( + moe._replica_plan_launch_args(plan) + ) + assert (compute_slots, helper_slots, generation, reserved_sms) == (11, 3, 7, 8) + assert actual_flags is ready_flags + moe._wait_rebalance_routes(plan) + wait_for_routes.assert_called_once_with() + + moe._wait_rebalance_routes() + wait_for_routes.assert_called_once_with() + + +@pytest.mark.cpu_only +def test_megamoe_replica_plan_rejects_unallocated_geometry(): + moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) + torch.nn.Module.__init__(moe) + moe.expert_size_per_partition = 8 + moe._rebalance_home_experts = 8 + moe._rebalance_slots_active = 3 + moe._quantize_reserved_sms = 8 + plan = MoEReplicaPlan( + resident_experts_per_rank=8, + compute_slots_per_rank=12, + ready_flags=torch.empty(8, dtype=torch.int64), + ready_generation=7, + reserved_sms=8, + _wait_for_routes=lambda: None, + ) + + with pytest.raises(RuntimeError, match="geometry does not match"): + moe._replica_plan_launch_args(plan) + + +@pytest.mark.cpu_only def test_megamoe_cutedsl_post_load_weights_uses_staged_hooks(): moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) torch.nn.Module.__init__(moe) quant_method = SimpleNamespace( transform_weights=MagicMock(), cache_derived_state=MagicMock(), + assert_live_weight_aliases=MagicMock(), ) moe.quant_method = quant_method + moe._rebalance_slots_active = 1 + arena = MagicMock() + # ``unittest.mock`` blocks dynamically-created ``assert_*`` attributes. + arena.assert_identity = MagicMock() + moe._rebalance_arena = arena + moe.layer_load_balancer = SimpleNamespace(bind_backend=MagicMock()) moe.post_load_weights() moe.transform_weights() quant_method.transform_weights.assert_called_once_with(moe) quant_method.cache_derived_state.assert_called_once_with(moe) + moe._rebalance_arena.assert_identity.assert_called_once_with() + quant_method.assert_live_weight_aliases.assert_called_once_with(moe, moe._rebalance_arena) + moe.layer_load_balancer.bind_backend.assert_called_once_with(moe) assert moe._weights_transformed is True @@ -2098,6 +2304,123 @@ def test_megamoe_cutedsl_mpi_bootstrap_keeps_executor_cuda_device(monkeypatch): ) +@pytest.mark.cpu_only +def test_megamoe_cutedsl_cache_derived_state_reuses_rebalance_binding( + monkeypatch, +): + from tensorrt_llm._torch.moe.fused_moe import PerLayerDynamicMoeLoadBalancer + from tensorrt_llm._torch.moe.fused_moe.mega_moe import rebalance_slot_scheduler + + config = SimpleNamespace( + num_slots=24, + auxiliary_sms=8, + validate_expert_capacity=lambda num_experts: None, + ) + manager = PerLayerDynamicMoeLoadBalancer( + ep_rank=0, + ep_size=4, + config=config, + ) + layer = manager.add_layer( + expert_count=16, + top_k=2, + slot_count_per_rank=6, + ) + arena = MagicMock(home_experts=4, helper_slots=2) + arena.assert_identity = MagicMock() + moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) + torch.nn.Module.__init__(moe) + moe.quant_method = SimpleNamespace( + cache_derived_state=MagicMock(), + assert_live_weight_aliases=MagicMock(), + ) + moe._rebalance_slots_active = 2 + moe._rebalance_arena = arena + moe.layer_load_balancer = layer + builder = MagicMock(return_value=SimpleNamespace(home_experts=4, helper_slots=2)) + monkeypatch.setattr( + rebalance_slot_scheduler, + "build_rebalance_slot_scheduler_group", + builder, + ) + + moe.cache_derived_state() + moe.cache_derived_state() + + assert moe.quant_method.cache_derived_state.call_count == 2 + assert arena.assert_identity.call_count == 2 + assert moe.quant_method.assert_live_weight_aliases.call_count == 2 + builder.assert_called_once_with(moe) + + +@pytest.mark.cpu_only +def test_rebalance_discard_plan_orders_release_and_allows_next_reuse(): + trace = [] + completion_event = object() + outputs = SimpleNamespace(physical_slot_ids=[3, 4]) + part = (outputs, 2) + copy_stream = object() + + group = RebalanceSlotSchedulerGroup.__new__(RebalanceSlotSchedulerGroup) + group._check_owner = MagicMock() + group._plan_part = part + group._generation = 1 + group._finished_generation = 0 + group._scheduled_generation = 1 + group._route_wait_generation = 0 + group._pending_release = None + group.plan_calls = 0 + group._consumer_done = completion_event + group._consumer_done_handle = 17 + group._execution_stream_handle = 23 + group.copy_stream = copy_stream + group.scheduler_stream = copy_stream + + def wait_for_routes(): + trace.append("wait_routes") + group._route_wait_generation = group._generation + + def record_event(event_handle, stream_handle): + assert (event_handle, stream_handle) == (17, 23) + trace.append("record_consumer_done") + + broadcaster = SimpleNamespace( + release_generation_after=MagicMock( + side_effect=lambda generation, event: trace.append("release_generation") + ), + mark_collective_reuse_safe=MagicMock( + side_effect=lambda authority, generation: trace.append("prove_reuse") + ), + ) + group.wait_for_routes = wait_for_routes + group._record_event = record_event + group.broadcaster = broadcaster + + group.discard_plan(part) + + assert group._plan_part is None + assert group.plan_calls == 1 + assert group._finished_generation == 1 + assert group._pending_release == (1, completion_event) + assert trace == ["wait_routes", "record_consumer_done", "release_generation"] + + # The next schedule advances this counter before proving that generation 1 + # can be reused. Exercise the real lease check to show cleanup left a valid + # transition rather than a permanently poisoned group. + group._scheduled_generation = 2 + lease = _LiveBankLeaseProvider(group) + lease.prove(1) + + assert lease.proved_generation == 1 + assert trace[-1] == "prove_reuse" + + +@pytest.mark.cpu_only +def test_rebalance_scheduler_group_requires_ep_process_group(): + parameters = inspect.signature(RebalanceSlotSchedulerGroup.__init__).parameters + assert "ep_process_group" in parameters + + def test_megamoe_load_weights_invalidates_cached_deepgemm_views(): method = W4A8MXFP4MXFP8MegaMoEDeepGemmMethod() hidden_size = 128 @@ -3354,29 +3677,32 @@ def test_enumerate_megamoe_candidate_tactics_curated_space( decode = megamoe_op.enumerate_megamoe_candidate_tactics(1024, sm_version=sm_version) prefill = megamoe_op.enumerate_megamoe_candidate_tactics(16384, sm_version=sm_version) assert (len(decode), len(prefill)) == (decode_count, prefill_count) - assert all(len(t) == 10 for t in decode + prefill) - assert {t[-1] for t in decode} == {(1, 1)} - assert {t[-1] for t in prefill} == {(2, 4)} + assert all(len(t) == 11 for t in decode + prefill) + assert {t[-2] for t in decode} == {(1, 1)} + assert {t[-2] for t in prefill} == {(2, 4)} + assert {t[-1] for t in decode + prefill} == {"expert"} for tactic in decode + prefill: megamoe_op.validate_megamoe_tactic(tactic, sm_version=sm_version) - for num_tokens in (64, 4096, 16384): - megamoe_op.validate_megamoe_tactic( - megamoe_op.default_megamoe_tactic(num_tokens), sm_version=sm_version - ) + default_tactics = [ + megamoe_op.default_megamoe_tactic(num_tokens) for num_tokens in (64, 4096, 16384) + ] + assert all(len(tactic) == 11 for tactic in default_tactics) + assert {tactic[-1] for tactic in default_tactics} == {"expert"} + for tactic in default_tactics: + megamoe_op.validate_megamoe_tactic(tactic, sm_version=sm_version) if sm_version == 107: for bucket, tactic in megamoe_op._SM107_GENPHASE_TACTICS.items(): megamoe_op.validate_megamoe_tactic(tactic, sm_version=sm_version) - assert ( - megamoe_op._default_megamoe_tactic_for_problem( - sm_version=sm_version, - max_tokens_per_rank=bucket, - num_tokens=bucket, - apply_topk_in_fc1=False, - in_kernel_fc2_reduce=False, - combine_format="bf16", - ) - == tactic + selected = megamoe_op._default_megamoe_tactic_for_problem( + sm_version=sm_version, + max_tokens_per_rank=bucket, + num_tokens=bucket, + apply_topk_in_fc1=False, + in_kernel_fc2_reduce=False, + combine_format="bf16", ) + assert len(selected) == 11 + assert selected == megamoe_op._unpack_tactic(tactic) invalid_tactic = list(megamoe_op.default_megamoe_tactic(64)) invalid_tactic[3] = ("grouped", 0) with pytest.raises(ValueError, match=r"schedule_policy hint must be a positive int or None"): @@ -3384,6 +3710,9 @@ def test_enumerate_megamoe_candidate_tactics_curated_space( legacy = ([256, 128, 256], [2, 1, 1], 512, "static", "epi_warps", True, 1, (1, 1)) megamoe_op.validate_megamoe_tactic(legacy, sm_version=sm_version) assert megamoe_op._unpack_tactic(legacy) == megamoe_op.default_megamoe_tactic(64) + v4 = megamoe_op.default_megamoe_tactic(64)[:-1] + megamoe_op.validate_megamoe_tactic(v4, sm_version=sm_version) + assert megamoe_op._unpack_tactic(v4) == megamoe_op.default_megamoe_tactic(64) @pytest.mark.gpu diff --git a/tests/unittest/_torch/moe/test_moe_impl.py b/tests/unittest/_torch/moe/test_moe_impl.py index d41abdc7ac4e..08fd8eb8a358 100644 --- a/tests/unittest/_torch/moe/test_moe_impl.py +++ b/tests/unittest/_torch/moe/test_moe_impl.py @@ -424,6 +424,35 @@ def test_pinned_megamoe_cutedsl_identity_resolves_to_the_leaf(): assert not report.degraded +def test_dynamic_eplb_is_a_declared_backend_capability(): + assert TrtllmCutedslMegaMoeNvfp4Impl.capabilities.supports_per_layer_dynamic_eplb + assert not DeepgemmCudaW4a8Mxfp4Mxfp8Impl.capabilities.supports_per_layer_dynamic_eplb + + +def test_per_layer_dynamic_eplb_filters_backends_by_declared_capability(): + deployment = dataclasses.replace( + _megamoe_cutedsl_deployment(), + eplb_enabled=True, + eplb_mode="per_layer_dynamic", + ) + accepted = resolve_moe_impl( + _megamoe_cutedsl_model_config(), + problem=_megamoe_cutedsl_problem(), + deployment=deployment, + impl_id=_MEGAMOE_CUTEDSL_IMPL_ID, + ) + assert impl_class_for(accepted) is TrtllmCutedslMegaMoeNvfp4Impl + + rejected = resolve_moe_impl( + _deepgemm_model_config(QuantAlgo.W4A8_MXFP4_MXFP8), + problem=_megamoe_cutedsl_problem(QuantAlgo.W4A8_MXFP4_MXFP8), + deployment=deployment, + impl_id="deepgemm.cuda.mega_moe.w4a8_mxfp4_mxfp8", + ) + assert rejected.winner is None + assert [item.reason for item in rejected.rejected] == [MoERejectReason.EPLB_UNSUPPORTED] + + def test_pinned_megamoe_cutedsl_identity_fails_hard_on_another_format(): """A pin that cannot run must fail, not widen to the sibling that can. diff --git a/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py b/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py new file mode 100644 index 000000000000..69fa28889b6f --- /dev/null +++ b/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py @@ -0,0 +1,411 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch + +from tensorrt_llm._torch.moe.fused_moe.mega_moe.rebalance_slot_scheduler import ( + RebalanceSlotSchedulerGroup, +) +from tensorrt_llm._torch.moe.fused_moe.moe_load_balancer import ( + MoeLoadBalancerIterContext, + get_moe_load_balancer, + maybe_create_moe_load_balancer, +) +from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb import PerLayerDynamicMoeLoadBalancer + +pytestmark = pytest.mark.cpu_only + + +class _Config: + mode = "per_layer_dynamic" + num_slots = 24 + auxiliary_sms = 8 + + def setup(self, ep_rank, ep_size): + self.ep_rank = ep_rank + self.ep_size = ep_size + + def validate_expert_capacity(self, num_experts): + if self.num_slots <= num_experts: + raise ValueError("num_slots must exceed num_experts") + + +class _StandardConfig: + mode = "standard" + layer_updates_per_iter = 0 + + def setup(self, ep_rank, ep_size): + raise AssertionError("unsupported standard EPLB must remain disabled") + + +class _Arena: + home_experts = 4 + helper_slots = 2 + + def __init__(self): + self.flags = torch.empty(4, dtype=torch.uint64) + self.generation = 0 + self.provider = SimpleNamespace(pool=None) + + def publish_ready_generation(self, generation): + self.generation = generation + + def terminal_flags_tensor(self): + return self.flags + + def assert_identity(self): + pass + + +class _Backend: + def __init__(self, arena): + self._rebalance_arena = arena + self.tactic_autotune = False + + +class _Group: + home_experts = 4 + helper_slots = 2 + + def __init__(self): + self.has_submission_owner = False + self.wait_calls = 0 + self.finish_calls = [] + self.handoff_calls = [] + self.release_calls = [] + self.events = [] + + def plan(self, routes, *, defer_wait): + assert defer_wait + self.has_submission_owner = True + self.events.append("plan") + return routes + 4, 7 + + def plan_schedule(self, routes): + self.has_submission_owner = True + self.events.append("schedule") + return routes + + def plan_finish(self, part, *, defer_wait): + assert defer_wait + self.events.append("finish") + return part + 4, 7 + + def discard_plan(self, part): + self.events.append("discard") + + def wait_for_routes(self): + self.wait_calls += 1 + + def finish(self, completion_event=None): + self.finish_calls.append(completion_event) + + def release_warmup_owner_after_sync(self, stream_handle): + self.handoff_calls.append(stream_handle) + self.has_submission_owner = False + + def release_quiesced_submission_owner(self, stream_handle): + self.release_calls.append(stream_handle) + self.has_submission_owner = False + + def close(self): + pass + + +def _manager(): + return PerLayerDynamicMoeLoadBalancer( + ep_rank=0, + ep_size=4, + config=_Config(), + ) + + +def test_layer_capacity_uses_configured_compute_slots(): + manager = _manager() + with pytest.raises(ValueError, match="configured local slot count"): + manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=4) + + +def test_layer_produces_and_finishes_replica_plan(): + manager = _manager() + layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) + arena = _Arena() + group = _Group() + layer.bind_runtime(arena=arena, scheduler_group=group) + + logical = torch.tensor([[0, 1], [2, 3]], dtype=torch.int32) + physical = layer.route(logical) + plan = layer.require_replica_plan() + + torch.testing.assert_close(physical, logical + 4) + assert plan.resident_experts_per_rank == 4 + assert plan.compute_slots_per_rank == 6 + assert plan.helper_slots_per_rank == 2 + assert plan.ready_generation == arena.generation == 7 + assert plan.ready_flags is arena.flags + assert plan.reserved_sms == 8 + plan.wait_for_routes() + assert group.wait_calls == 1 + + layer.finish_replica_plan(plan) + assert group.finish_calls == [None] + assert layer.replica_plan is None + + +def test_gap_hook_runs_between_schedule_and_plan_finish(): + manager = _manager() + layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) + arena = _Arena() + group = _Group() + layer.bind_runtime(arena=arena, scheduler_group=group) + + logical = torch.tensor([[0, 1]], dtype=torch.int32) + physical = layer.route(logical, gap_hook=lambda: group.events.append("hook")) + torch.testing.assert_close(physical, logical + 4) + assert group.events == ["schedule", "hook", "finish"] + layer.finish_replica_plan() + + +def test_gap_hook_failure_discards_generation_and_adds_context(): + manager = _manager() + layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) + group = _Group() + layer.bind_runtime(arena=_Arena(), scheduler_group=group) + + def fail(): + raise RuntimeError("shared expert failed") + + with pytest.raises(RuntimeError, match="shared expert failed") as caught: + layer.route(torch.tensor([[0, 1]], dtype=torch.int32), gap_hook=fail) + assert group.events == ["schedule", "discard"] + assert "safely released" in caught.value.__notes__[0] + assert layer.replica_plan is None + + +def test_layer_disables_standard_host_migration_and_checks_geometry(): + manager = _manager() + layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) + assert not layer.is_dynamic_routing() + assert not layer.need_load_shared_weights() + assert layer.get_load_expert_ids() == [0, 1, 2, 3] + assert layer.get_local_statistic_tensor() is None + layer.set_initial_weight_assignments(range(16)) + with pytest.raises(ValueError, match="custom initial_global_assignments"): + layer.set_initial_weight_assignments([0]) + with pytest.raises(ValueError, match="arena H/S geometry"): + layer.bind_runtime( + arena=SimpleNamespace(home_experts=5, helper_slots=1), + scheduler_group=_Group(), + ) + + +def test_bind_backend_is_idempotent_before_scheduler_construction(monkeypatch): + from tensorrt_llm._torch.moe.fused_moe.mega_moe import rebalance_slot_scheduler + + manager = _manager() + layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) + arena = _Arena() + backend = _Backend(arena) + group = _Group() + builder = MagicMock(return_value=group) + monkeypatch.setattr( + rebalance_slot_scheduler, + "build_rebalance_slot_scheduler_group", + builder, + ) + + layer.bind_backend(backend) + layer.bind_backend(backend) + + builder.assert_called_once_with(backend) + assert layer._arena is arena + assert layer._scheduler_group is group + + +def test_bind_backend_rejects_different_backend_or_arena_before_construction( + monkeypatch, +): + from tensorrt_llm._torch.moe.fused_moe.mega_moe import rebalance_slot_scheduler + + manager = _manager() + layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) + arena = _Arena() + backend = _Backend(arena) + builder = MagicMock(return_value=_Group()) + monkeypatch.setattr( + rebalance_slot_scheduler, + "build_rebalance_slot_scheduler_group", + builder, + ) + layer.bind_backend(backend) + + with pytest.raises(RuntimeError, match="different backend"): + layer.bind_backend(_Backend(arena)) + backend._rebalance_arena = _Arena() + with pytest.raises(RuntimeError, match="changed its rebalance arena"): + layer.bind_backend(backend) + + builder.assert_called_once_with(backend) + + +def test_iteration_context_reuses_common_lifecycle(): + manager = _manager() + with manager: + assert get_moe_load_balancer() is manager + with MoeLoadBalancerIterContext(manager, True, True): + assert manager.in_iter + assert not manager.in_iter + assert manager.iter_id == 1 + assert get_moe_load_balancer() is None + + +def test_public_warmup_handoff_and_model_owned_shutdown(): + manager = _manager() + layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) + arena = _Arena() + pool = MagicMock() + pool.close_collectively.return_value = True + arena.provider.pool = pool + manager._rebalance_shared_slot_pools[("test",)] = pool + manager._rebalance_shared_slot_pool_order.append(pool) + group = _Group() + group.has_submission_owner = True + layer.bind_runtime(arena=arena, scheduler_group=group) + + assert manager.has_submission_owner + manager.handoff_submission_owner(1234) + assert group.handoff_calls == [1234] + manager.release_quiesced_submission_owner(1234) + assert group.release_calls == [1234] + assert manager.shutdown() + pool.close_collectively.assert_called_once_with(local_safe=True) + assert manager.is_shutdown + + +def _scheduler_group_for_owner_release(*, finished=True): + group = RebalanceSlotSchedulerGroup.__new__(RebalanceSlotSchedulerGroup) + group._closed = False + group._plan_part = None + group._generation = 3 + group._finished_generation = 3 if finished else 2 + group.plan_calls = 3 + group._owner_thread_id = 99 + group._execution_stream_handle = 1234 + return group + + +def test_quiesced_owner_release_clears_worker_and_stream_binding(): + group = _scheduler_group_for_owner_release() + + group.release_quiesced_submission_owner(1234) + + assert group._owner_thread_id is None + assert group._execution_stream_handle is None + + +def test_quiesced_owner_release_rejects_unfinished_generation(): + group = _scheduler_group_for_owner_release(finished=False) + + with pytest.raises(RuntimeError, match="all plans to be finished"): + group.release_quiesced_submission_owner(1234) + + +def test_ep_communicator_is_model_owned(monkeypatch): + from tensorrt_llm._torch.moe.fused_moe.mega_moe import rebalance_live_arena + + creations = [] + + class Comm: + def __init__(self, process_group): + self.process_group = process_group + self.validations = [] + creations.append(self) + + def validate_geometry(self, *, expected_rank, expected_world): + self.validations.append((expected_rank, expected_world)) + + monkeypatch.setattr(rebalance_live_arena, "_EpComm", Comm) + manager = _manager() + process_group = object() + + first = manager.get_ep_comm(process_group) + second = manager.get_ep_comm(process_group) + + assert first is second + assert len(creations) == 1 + assert first.validations == [(0, 4), (0, 4)] + with pytest.raises(RuntimeError, match="cannot span EP ProcessGroups"): + manager.get_ep_comm(object()) + + +def test_factory_selects_per_layer_dynamic_manager(): + mapping = SimpleNamespace( + moe_ep_rank=0, + moe_ep_size=4, + moe_cluster_size=1, + ) + config = _Config() + model_config = SimpleNamespace( + mapping=mapping, + pretrained_config=SimpleNamespace(architectures=["DeepseekV4ForCausalLM"]), + moe_load_balancer=config, + ) + + with maybe_create_moe_load_balancer(model_config, mapping) as manager: + assert isinstance(manager, PerLayerDynamicMoeLoadBalancer) + assert config.ep_rank == 0 + assert config.ep_size == 4 + + +@pytest.mark.parametrize( + "architecture, ep_size, cluster_size, message", + [ + ("UnsupportedForCausalLM", 4, 1, "does not support model architecture"), + ("DeepseekV4ForCausalLM", 1, 1, "requires expert parallelism"), + ("DeepseekV4ForCausalLM", 4, 2, "incompatible with smart routing"), + ], +) +def test_per_layer_dynamic_factory_fails_fast_for_unsupported_topology( + architecture, ep_size, cluster_size, message +): + mapping = SimpleNamespace( + moe_ep_rank=0, + moe_ep_size=ep_size, + moe_cluster_size=cluster_size, + ) + model_config = SimpleNamespace( + mapping=mapping, + pretrained_config=SimpleNamespace(architectures=[architecture]), + moe_load_balancer=_Config(), + ) + + with pytest.raises(ValueError, match=message): + maybe_create_moe_load_balancer(model_config, mapping) + + +@pytest.mark.parametrize( + "architecture, ep_size, cluster_size", + [ + ("UnsupportedForCausalLM", 4, 1), + ("DeepseekV4ForCausalLM", 1, 1), + ("DeepseekV4ForCausalLM", 4, 2), + ], +) +def test_standard_factory_keeps_unsupported_topology_disabled(architecture, ep_size, cluster_size): + mapping = SimpleNamespace( + moe_ep_rank=0, + moe_ep_size=ep_size, + moe_cluster_size=cluster_size, + ) + model_config = SimpleNamespace( + mapping=mapping, + pretrained_config=SimpleNamespace(architectures=[architecture]), + moe_load_balancer=_StandardConfig(), + ) + + with maybe_create_moe_load_balancer(model_config, mapping) as manager: + assert manager is None diff --git a/tests/unittest/_torch/weight_sharing/_source_identity_fakes.py b/tests/unittest/_torch/weight_sharing/_source_identity_fakes.py index e3c31c33712f..b992d39812f1 100644 --- a/tests/unittest/_torch/weight_sharing/_source_identity_fakes.py +++ b/tests/unittest/_torch/weight_sharing/_source_identity_fakes.py @@ -151,6 +151,7 @@ def __init__( quant_config=_UNSET, attn_backend: str = "TRTLLM", moe_backend: str = "CUTLASS", + moe_load_balancer=None, ): self.mapping = mapping or FakeMapping() self.pretrained_config = ( @@ -161,6 +162,7 @@ def __init__( self.force_dynamic_quantization = False self.attn_backend = attn_backend self.moe_backend = moe_backend + self.moe_load_balancer = moe_load_balancer self.nvfp4_gemm_allowed_backends = ["cutlass", "cublaslt", "cuda_core"] self.moe_disable_finalize_fusion = False self.use_low_precision_moe_combine = False diff --git a/tests/unittest/_torch/weight_sharing/test_source_identity.py b/tests/unittest/_torch/weight_sharing/test_source_identity.py index 8ee4acd7b929..3d7ab10b2c37 100644 --- a/tests/unittest/_torch/weight_sharing/test_source_identity.py +++ b/tests/unittest/_torch/weight_sharing/test_source_identity.py @@ -43,6 +43,7 @@ SourceIdentityMismatchError, check_weight_sharing_compatibility, ) +from tensorrt_llm.llmapi.llm_args import MoeLoadBalancerConfig, MoePerLayerDynamicLoadBalancerConfig pytestmark = pytest.mark.cpu_only @@ -132,6 +133,21 @@ def test_backend_mismatch_flags_global(): assert "backend_fingerprint" in result.mismatched_fields +def test_moe_load_balancer_config_controls_backend_identity(): + base = identity_from(FakeModelConfig()) + standard = identity_from(FakeModelConfig(moe_load_balancer=MoeLoadBalancerConfig(num_slots=64))) + per_layer_dynamic = identity_from( + FakeModelConfig(moe_load_balancer=MoePerLayerDynamicLoadBalancerConfig(num_slots=72)) + ) + same_per_layer_dynamic = identity_from( + FakeModelConfig(moe_load_balancer=MoePerLayerDynamicLoadBalancerConfig(num_slots=72)) + ) + + assert not base.matches(standard).matched + assert not standard.matches(per_layer_dynamic).matched + assert per_layer_dynamic.matches(same_per_layer_dynamic).matched + + def test_model_layout_field_mismatch_flags_global_and_shard(): a = identity_from(FakeModelConfig(pretrained_config=FakePretrainedConfig(hidden_size=4096))) b = identity_from(FakeModelConfig(pretrained_config=FakePretrainedConfig(hidden_size=8192))) diff --git a/tests/unittest/executor/test_shutdown_retry_contracts.py b/tests/unittest/executor/test_shutdown_retry_contracts.py new file mode 100644 index 000000000000..f446c559a29d --- /dev/null +++ b/tests/unittest/executor/test_shutdown_retry_contracts.py @@ -0,0 +1,584 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only lifecycle contracts for retryable Ray and RPC shutdown.""" + +import ast +import sys +import threading +import types +from pathlib import Path +from types import SimpleNamespace + +import pytest + +pytestmark = pytest.mark.cpu_only + +_ROOT = Path(__file__).resolve().parents[3] + + +def _extract_function(source, function_name, namespace): + tree = ast.parse(source.read_text()) + function = next( + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == function_name + ) + module = ast.fix_missing_locations(ast.Module(body=[function], type_ignores=[])) + exec(compile(module, str(source), "exec"), namespace) + return namespace[function_name] + + +def _extract_class(source, class_name, methods, namespace, bases=()): + tree = ast.parse(source.read_text()) + original = next( + node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == class_name + ) + body = [ + node + for node in original.body + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name in methods + ] + derived = ast.ClassDef( + name=class_name, + bases=[ast.Name(id=base, ctx=ast.Load()) for base in bases], + keywords=[], + body=body, + decorator_list=[], + ) + module = ast.fix_missing_locations(ast.Module(body=[derived], type_ignores=[])) + exec(compile(module, str(source), "exec"), namespace) + return namespace[class_name] + + +class _Timeout(Exception): + pass + + +class _Ref: + def __init__(self, outcomes): + self.outcomes = list(outcomes) + + def get(self): + outcome = self.outcomes.pop(0) if len(self.outcomes) > 1 else self.outcomes[0] + if isinstance(outcome, BaseException): + raise outcome + return outcome + + +class _RemoteShutdown: + def __init__(self, refs): + self.refs = list(refs) + self.calls = 0 + + def remote(self): + ref = self.refs[min(self.calls, len(self.refs) - 1)] + self.calls += 1 + return ref + + +class _RayWorker: + def __init__(self, refs): + self.shutdown = _RemoteShutdown(refs) + + +class _FakeRay: + def __init__(self): + self.exceptions = SimpleNamespace(GetTimeoutError=_Timeout) + self.util = SimpleNamespace(remove_placement_group=lambda group: None) + self.killed = [] + + @staticmethod + def get(ref, timeout=None): + return ref.get() + + def kill(self, worker, no_restart=True): + self.killed.append(worker) + + @staticmethod + def is_initialized(): + return False + + @staticmethod + def shutdown(): + return None + + +def _ray_executor(ray): + cls = _extract_class( + _ROOT / "tensorrt_llm/executor/ray/executor.py", + "RayExecutor", + {"shutdown"}, + { + "ray": ray, + "time": __import__("time"), + "logger": SimpleNamespace( + warning=lambda *args, **kwargs: None, debug=lambda *args, **kwargs: None + ), + "logger_debug": lambda *args, **kwargs: None, + }, + ) + executor = cls() + executor._shutdown_event = threading.Event() + executor.main_loop = None + executor.main_loop_task_obj = None + executor.rpc_client = SimpleNamespace(close=lambda: None) + executor.placement_group = None + executor.bundle_indices = None + executor.has_start_local_cluser = False + executor._wait_for_cluster_resource_release = lambda timeout: None + return executor + + +def test_ray_shutdown_failure_preserves_actors_for_retry(): + ray = _FakeRay() + worker = _RayWorker([_Ref([RuntimeError("cleanup failed")]), _Ref([None])]) + executor = _ray_executor(ray) + executor.workers = [worker] + + with pytest.raises(RuntimeError, match="cleanup failed"): + executor.shutdown() + + assert executor.workers == [worker] + assert not executor._shutdown_event.is_set() + assert ray.killed == [] + + executor.shutdown() + assert worker.shutdown.calls == 2 + assert ray.killed == [worker] + assert executor.workers is None + + +def test_ray_shutdown_timeout_reuses_pending_ref(): + ray = _FakeRay() + pending = _Ref([_Timeout("still running"), None]) + worker = _RayWorker([pending]) + executor = _ray_executor(ray) + executor.workers = [worker] + + with pytest.raises(_Timeout, match="still running"): + executor.shutdown() + assert executor._pending_shutdown_refs == [pending] + + executor.shutdown() + assert worker.shutdown.calls == 1 + assert executor.workers is None + + +class _BaseWorker: + def shutdown(self): + outcome = self.base_outcomes.pop(0) + if isinstance(outcome, BaseException): + raise outcome + + +class _Comm: + def __init__(self, gathered): + self.gathered = list(gathered) + self.broadcasts = [] + + def allgather(self, value): + return self.gathered.pop(0) + + def bcast(self, value, root): + self.broadcasts.append(value) + return value + + +def _rpc_worker(comm): + cls = _extract_class( + _ROOT / "tensorrt_llm/executor/rpc_worker.py", + "RpcWorker", + {"shutdown", "commit_shutdown", "_complete_shutdown_commit"}, + { + "Base": _BaseWorker, + "mpi_rank": lambda: 0, + "mpi_comm": lambda: comm, + "logger_debug": lambda *args, **kwargs: None, + "_RPC_SHUTDOWN_RETRY": "retry", + "_RPC_SHUTDOWN_COMMIT": "commit", + }, + bases=("Base",), + ) + worker = cls() + worker._rpc_collective_shutdown = True + worker._rpc_shutdown_attempted = False + worker._rpc_shutdown_ready = False + worker._rpc_shutdown_committed = False + worker.shutdown_event = threading.Event() + return worker + + +def test_rpc_cleanup_retry_and_commit_are_rank_aligned_and_idempotent(): + comm = _Comm([["RuntimeError: local failure", None], [None, None]]) + worker = _rpc_worker(comm) + worker.base_outcomes = [RuntimeError("local failure"), None] + + with pytest.raises(RuntimeError, match="local failure"): + worker.shutdown() + assert not worker._rpc_shutdown_ready + assert comm.broadcasts == [] + + worker.shutdown() + assert worker._rpc_shutdown_ready + assert comm.broadcasts == ["retry"] + + worker.commit_shutdown() + worker.commit_shutdown() + assert comm.broadcasts == ["retry", "commit"] + assert not worker.shutdown_event.is_set() + + worker._complete_shutdown_commit() + assert worker.shutdown_event.is_set() + + +def test_rpc_peer_failure_is_propagated_to_rank_zero(): + comm = _Comm([[None, "RuntimeError: peer failure"]]) + worker = _rpc_worker(comm) + worker.base_outcomes = [None] + + with pytest.raises(RuntimeError, match="rank 1: RuntimeError: peer failure"): + worker.shutdown() + assert not worker._rpc_shutdown_ready + + +class _RemoteCall: + def __init__(self, name, calls, error=None): + self.name = name + self.calls = calls + self.error = error + + def remote(self, *, need_response): + self.calls.append((self.name, need_response)) + if self.error is not None: + raise self.error + + +class _RpcClient: + def __init__(self, calls, shutdown_error=None): + self.calls = calls + self.shutdown_error = shutdown_error + + def shutdown(self): + return _RemoteCall("cleanup", self.calls, self.shutdown_error) + + def commit_shutdown(self): + return _RemoteCall("commit", self.calls) + + def close(self): + self.calls.append(("close", True)) + + +def _rpc_proxy(): + return _extract_class( + _ROOT / "tensorrt_llm/executor/rpc_proxy.py", + "GenerationExecutorRpcProxy", + {"shutdown_remote", "shutdown"}, + { + "logger_debug": lambda *args, **kwargs: None, + "logger": SimpleNamespace(warning=lambda *args, **kwargs: None), + "threading": threading, + }, + ) + + +def test_rpc_proxy_requires_cleanup_and_commit_responses(): + calls = [] + proxy = _rpc_proxy()() + proxy.rpc_client = _RpcClient(calls) + + proxy.shutdown_remote() + + assert calls == [("cleanup", True), ("commit", True)] + + +def test_rpc_proxy_cleanup_failure_keeps_retry_gate_and_resources_live(): + calls = [] + proxy = _rpc_proxy()() + proxy._shutdown_event = threading.Event() + proxy.rpc_client = _RpcClient(calls, RuntimeError("remote cleanup failed")) + proxy.main_loop = None + proxy.main_loop_task_obj = None + proxy.main_loop_thread = None + proxy.mpi_session = object() + + with pytest.raises(RuntimeError, match="remote cleanup failed"): + proxy.shutdown() + + assert not proxy._shutdown_event.is_set() + assert proxy.mpi_session is not None + assert calls == [("cleanup", True)] + + +def test_rpc_server_completes_commit_only_after_response_send(): + source = _ROOT / "tensorrt_llm/executor/rpc/rpc_server.py" + tree = ast.parse(source.read_text()) + server = next( + node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "RPCServer" + ) + process = next( + node + for node in server.body + if isinstance(node, ast.AsyncFunctionDef) and node.name == "_process_requests" + ) + send_line = min( + node.lineno + for node in ast.walk(process) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "_send_response" + ) + complete_line = min( + node.lineno + for node in ast.walk(process) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "complete_commit" + ) + method_source = ast.get_source_segment(source.read_text(), process) + + assert send_line < complete_line + assert 'req.method_name == "commit_shutdown"' in method_source + assert "response.error is None" in method_source + + +def test_symm_buffer_release_retains_failures_for_retry(): + cache = { + "failed": ("bad", object()), + "released": ("good", object()), + } + attempts = {"bad": 0} + + def free(buffer): + if buffer == "bad" and attempts["bad"] == 0: + attempts["bad"] += 1 + raise RuntimeError("destroy failed") + return 16 + + release = _extract_function( + _ROOT / "tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py", + "release_symm_buffer_cache", + { + "_MEGA_MOE_SYMM_BUFFER_CACHE": cache, + "_free_symm_buffer": free, + "logger": SimpleNamespace( + error=lambda *args, **kwargs: None, info=lambda *args, **kwargs: None + ), + }, + ) + + with pytest.raises(RuntimeError, match="failed to release 1"): + release() + assert list(cache) == ["failed"] + + release() + assert cache == {} + + +def test_mpi_worker_retries_symm_release_before_destroying_process_group(): + events = [] + module_name = "test_retryable_megamoe_symm_release" + release_attempts = 0 + + def release_symm_buffer_cache(): + nonlocal release_attempts + release_attempts += 1 + events.append("release_symm") + if release_attempts == 1: + raise RuntimeError("symm release failed") + + mega_moe = types.ModuleType(module_name) + mega_moe.release_symm_buffer_cache = release_symm_buffer_cache + fake_distributed = types.ModuleType("torch.distributed") + fake_distributed.is_initialized = lambda: True + fake_distributed.destroy_process_group = lambda: events.append("destroy_pg") + fake_torch = types.ModuleType("torch") + fake_torch.distributed = fake_distributed + fake_torch.cuda = SimpleNamespace(empty_cache=lambda: events.append("empty_cache")) + + previous = {name: sys.modules.get(name) for name in (module_name, "torch", "torch.distributed")} + sys.modules[module_name] = mega_moe + sys.modules["torch"] = fake_torch + sys.modules["torch.distributed"] = fake_distributed + try: + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/worker.py", + "GenerationExecutorWorker", + {"shutdown"}, + { + "logger_debug": lambda *args, **kwargs: None, + "mpi_rank": lambda: 0, + "mpi_comm": lambda: SimpleNamespace(allgather=lambda value: [value]), + "sys": sys, + "_MEGA_MOE_DEEPGEMM_MODULE": module_name, + "gc": SimpleNamespace(collect=lambda: events.append("gc")), + }, + ) + engine = SimpleNamespace( + shutdown_all_ranks=True, + can_enqueue_requests=lambda: False, + shutdown=lambda: events.append("engine_shutdown"), + terminal_cleanup=lambda *, local_safe: events.append(("terminal_cleanup", local_safe)), + ) + worker = worker_cls() + worker.doing_shutdown = False + worker.engine = engine + worker.llm_args = SimpleNamespace(backend="pytorch") + worker._handle_background_error = lambda: events.append("background") + + with pytest.raises(RuntimeError, match="symm release failed"): + worker.shutdown() + assert events == [ + "engine_shutdown", + ("terminal_cleanup", True), + "release_symm", + ] + assert worker.engine is engine + assert not worker.doing_shutdown + + worker.shutdown() + assert events == [ + "engine_shutdown", + ("terminal_cleanup", True), + "release_symm", + ("terminal_cleanup", True), + "release_symm", + "destroy_pg", + "gc", + "empty_cache", + "background", + ] + assert worker.engine is None + finally: + for name, module in previous.items(): + if module is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = module + + +def test_mpi_worker_peer_cleanup_failure_blocks_later_collectives(): + events = [] + outcomes = [ + [None, "RuntimeError: peer cleanup failed"], + [None, None], + [None, None], + ] + + class Comm: + def allgather(self, value): + events.append(("allgather", value)) + return outcomes.pop(0) + + fake_distributed = types.ModuleType("torch.distributed") + fake_distributed.is_initialized = lambda: True + fake_distributed.destroy_process_group = lambda: events.append("destroy_pg") + fake_torch = types.ModuleType("torch") + fake_torch.distributed = fake_distributed + fake_torch.cuda = SimpleNamespace(empty_cache=lambda: events.append("empty_cache")) + + previous = {name: sys.modules.get(name) for name in ("torch", "torch.distributed")} + sys.modules["torch"] = fake_torch + sys.modules["torch.distributed"] = fake_distributed + try: + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/worker.py", + "GenerationExecutorWorker", + {"shutdown"}, + { + "logger_debug": lambda *args, **kwargs: None, + "mpi_rank": lambda: 0, + "mpi_comm": lambda: Comm(), + "sys": sys, + "_MEGA_MOE_DEEPGEMM_MODULE": "missing_peer_failure_megamoe", + "gc": SimpleNamespace(collect=lambda: events.append("gc")), + }, + ) + engine = SimpleNamespace( + shutdown_all_ranks=True, + can_enqueue_requests=lambda: False, + shutdown=lambda: events.append("engine_shutdown"), + terminal_cleanup=lambda *, local_safe: events.append(("terminal_cleanup", local_safe)), + ) + worker = worker_cls() + worker.doing_shutdown = False + worker.engine = engine + worker.llm_args = SimpleNamespace(backend="pytorch") + worker._handle_background_error = lambda: events.append("background") + + with pytest.raises(RuntimeError, match="rank 1.*peer cleanup failed"): + worker.shutdown() + assert events == [ + "engine_shutdown", + ("terminal_cleanup", True), + ("allgather", None), + ] + assert worker.engine is engine + assert not worker.doing_shutdown + + worker.shutdown() + assert events == [ + "engine_shutdown", + ("terminal_cleanup", True), + ("allgather", None), + ("terminal_cleanup", True), + ("allgather", None), + ("allgather", None), + "destroy_pg", + "gc", + "empty_cache", + "background", + ] + assert outcomes == [] + assert worker.engine is None + finally: + for name, module in previous.items(): + if module is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = module + + +def test_mpi_worker_background_error_keeps_shutdown_committed(monkeypatch): + cleanups = [] + fake_distributed = types.ModuleType("torch.distributed") + fake_distributed.is_initialized = lambda: True + fake_distributed.destroy_process_group = lambda: None + fake_torch = types.ModuleType("torch") + fake_torch.distributed = fake_distributed + fake_torch.cuda = SimpleNamespace(empty_cache=lambda: None) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + monkeypatch.setitem(sys.modules, "torch.distributed", fake_distributed) + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/worker.py", + "GenerationExecutorWorker", + {"shutdown"}, + { + "logger_debug": lambda *args, **kwargs: None, + "mpi_rank": lambda: 0, + "mpi_comm": lambda: SimpleNamespace(allgather=lambda value: [value]), + "sys": sys, + "_MEGA_MOE_DEEPGEMM_MODULE": "missing_background_error_megamoe", + "gc": SimpleNamespace(collect=lambda: None), + }, + ) + worker = worker_cls() + worker.doing_shutdown = False + worker.engine = SimpleNamespace( + shutdown_all_ranks=True, + can_enqueue_requests=lambda: False, + shutdown=lambda: None, + terminal_cleanup=lambda *, local_safe: cleanups.append(local_safe), + ) + worker.llm_args = SimpleNamespace(backend="pytorch") + + def background_error(): + raise RuntimeError("background") + + worker._handle_background_error = background_error + + with pytest.raises(RuntimeError, match="background"): + worker.shutdown() + # Peers have committed teardown; a retry must not re-enter its collectives. + assert worker.doing_shutdown and worker.engine is None + worker.shutdown() + assert cleanups == [True] diff --git a/tests/unittest/llmapi/test_gms_args.py b/tests/unittest/llmapi/test_gms_args.py index d2f3e162b12e..dd8cf5f51a9a 100644 --- a/tests/unittest/llmapi/test_gms_args.py +++ b/tests/unittest/llmapi/test_gms_args.py @@ -127,6 +127,51 @@ def test_non_default_config_warns_when_gms_inactive( assert relevant == [] +class TestGmsMoeCompatibility: + @pytest.mark.parametrize( + "load_balancer", + [ + pytest.param({"num_slots": 64}, id="standard"), + pytest.param( + {"mode": "per_layer_dynamic", "num_slots": 72}, + id="per-layer-dynamic", + ), + ], + ) + def test_rejects_gms_with_load_balancer(self, load_balancer): + with pytest.raises( + ValueError, + match=r"LoadFormat\.GMS.*moe_config\.load_balancer", + ): + _make_args( + load_format=LoadFormat.GMS, + moe_config={"load_balancer": load_balancer}, + ) + + @pytest.mark.parametrize( + "load_format, load_balancer", + [ + pytest.param(LoadFormat.GMS, None, id="gms-without-eplb"), + pytest.param( + LoadFormat.AUTO, + {"num_slots": 64}, + id="standard-without-gms", + ), + pytest.param( + LoadFormat.AUTO, + {"mode": "per_layer_dynamic", "num_slots": 72}, + id="per-layer-dynamic-without-gms", + ), + ], + ) + def test_accepts_compatible_load_balancer(self, load_format, load_balancer): + args = _make_args( + load_format=load_format, + moe_config={"load_balancer": load_balancer}, + ) + assert args.load_format == load_format + + class TestLoadFormatGms: def test_gms_enum_present(self): assert LoadFormat.GMS.name == "GMS" diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 237e6a3183b3..bcaf206da232 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -54,6 +54,8 @@ EncodeCudaGraphConfig, ExecutorMemoryType, KvCacheConfig, MambaStateConfig, MoeConfig, + MoeLoadBalancerConfig, + MoePerLayerDynamicLoadBalancerConfig, MTPDecodingConfig, MultimodalConfig, MultimodalEncoderCudaGraphConfig, NGramDecodingConfig, PeftCacheConfig, @@ -249,6 +251,205 @@ def test_ignores_override_without_neural_drafter(self) -> None: assert llm_args.speculative_config.moe_backend == "CUTLASS" +@pytest.mark.cpu_only +class TestMoeConfigPerLayerDynamicEplb: + + @staticmethod + def _per_layer_dynamic() -> dict[str, object]: + return {"mode": "per_layer_dynamic", "num_slots": 72} + + @pytest.mark.parametrize("backend", ["CUTLASS", "DEEPGEMM"]) + def test_rejects_explicit_incompatible_backend(self, backend: str) -> None: + with pytest.raises(ValidationError, + match=r"EPLB requires backend='MEGAMOE_CUTEDSL'"): + MoeConfig(backend=backend, load_balancer=self._per_layer_dynamic()) + + @pytest.mark.parametrize("backend", [None, "AUTO", "MEGAMOE_CUTEDSL"]) + def test_resolves_auto_and_accepts_compatible_backend( + self, backend: str | None) -> None: + kwargs = {} if backend is None else {"backend": backend} + + config = MoeConfig(load_balancer=self._per_layer_dynamic(), **kwargs) + + assert isinstance(config.load_balancer, + MoePerLayerDynamicLoadBalancerConfig) + assert config.backend == "MEGAMOE_CUTEDSL" + assert config.load_balancer.auxiliary_sms == 8 + + @pytest.mark.parametrize( + "num_experts, message", + [ + pytest.param(0, "must be positive", id="zero-experts"), + pytest.param(65, "must be divisible", id="non-divisible-experts"), + pytest.param(72, "must exceed", id="no-helper-capacity"), + pytest.param(80, "must exceed", id="negative-helper-capacity"), + ], + ) + def test_validates_physical_slot_capacity_after_setup( + self, num_experts: int, message: str) -> None: + config = MoePerLayerDynamicLoadBalancerConfig(num_slots=72) + config.setup(ep_rank=0, ep_size=8) + + with pytest.raises(ValueError, match=message): + config.validate_expert_capacity(num_experts) + + def test_accepts_positive_helper_capacity_after_setup(self) -> None: + config = MoePerLayerDynamicLoadBalancerConfig(num_slots=72) + config.setup(ep_rank=0, ep_size=8) + + config.validate_expert_capacity(num_experts=64) + assert config.num_local_slots == 9 + + def test_expert_capacity_requires_setup(self) -> None: + config = MoePerLayerDynamicLoadBalancerConfig(num_slots=72) + with pytest.raises(ValueError, match=r"Call setup\(\)"): + config.validate_expert_capacity(num_experts=64) + + def test_legacy_dictionary_defaults_to_standard_mode(self) -> None: + config = MoeConfig( + backend="CUTLASS", + load_balancer={ + "num_slots": 64, + "layer_updates_per_iter": 2 + }, + ) + + assert isinstance(config.load_balancer, MoeLoadBalancerConfig) + assert config.load_balancer.mode == "standard" + assert config.load_balancer.layer_updates_per_iter == 2 + assert config.backend == "CUTLASS" + + @pytest.mark.parametrize("field", ["num_slots", "auxiliary_sms"]) + def test_per_layer_dynamic_positive_fields(self, field: str) -> None: + values = self._per_layer_dynamic() + values[field] = 0 + with pytest.raises(ValidationError, match=field): + MoeConfig(load_balancer=values) + + def test_per_layer_dynamic_requires_eight_auxiliary_sms(self) -> None: + values = self._per_layer_dynamic() + values["auxiliary_sms"] = 7 + with pytest.raises(ValidationError, match="auxiliary_sms"): + MoeConfig(load_balancer=values) + + def test_per_layer_dynamic_accepts_larger_auxiliary_sm_budget(self) -> None: + values = self._per_layer_dynamic() + values["auxiliary_sms"] = 12 + config = MoeConfig(load_balancer=values) + assert config.load_balancer.auxiliary_sms == 12 + + def test_per_layer_dynamic_rejects_more_than_halo_q_cta_limit(self) -> None: + values = self._per_layer_dynamic() + values["auxiliary_sms"] = 129 + with pytest.raises(ValidationError, match="auxiliary_sms"): + MoeConfig(load_balancer=values) + + def test_per_layer_dynamic_disables_decode_cuda_graph(self) -> None: + args = TorchLlmArgs( + model=llama_model_path, + moe_config={"load_balancer": self._per_layer_dynamic()}) + assert args.cuda_graph_config is None + + def test_per_layer_dynamic_rejects_prefill_cuda_graph(self) -> None: + with pytest.raises(ValidationError, match="prefill CUDA graphs"): + TorchLlmArgs( + model=llama_model_path, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.PIECEWISE, + moe_config={"load_balancer": self._per_layer_dynamic()}) + + def test_legacy_yaml_defaults_to_standard_mode(self, + tmp_path: Path) -> None: + config_path = tmp_path / "eplb.yaml" + config_path.write_text("num_slots: 64\nlayer_updates_per_iter: 2\n") + + args = TorchLlmArgs( + model=llama_model_path, + moe_config={"load_balancer": str(config_path)}, + ) + + assert isinstance(args.moe_config.load_balancer, MoeLoadBalancerConfig) + assert args.moe_config.load_balancer.mode == "standard" + + def test_per_layer_dynamic_yaml_sets_backend(self, tmp_path: Path) -> None: + config_path = tmp_path / "eplb.yaml" + config_path.write_text("mode: per_layer_dynamic\nnum_slots: 72\n") + + args = TorchLlmArgs( + model=llama_model_path, + moe_config={"load_balancer": str(config_path)}, + ) + + assert isinstance(args.moe_config.load_balancer, + MoePerLayerDynamicLoadBalancerConfig) + assert args.moe_config.backend == "MEGAMOE_CUTEDSL" + + def test_per_layer_dynamic_yaml_is_typed_by_moe_config_before_llm_args( + self, tmp_path: Path) -> None: + config_path = tmp_path / "eplb.yaml" + config_path.write_text("mode: per_layer_dynamic\nnum_slots: 72\n") + + config = MoeConfig(load_balancer=str(config_path)) + + assert isinstance(config.load_balancer, + MoePerLayerDynamicLoadBalancerConfig) + assert config.backend == "MEGAMOE_CUTEDSL" + + def test_per_layer_dynamic_yaml_configures_queue_before_cuda_init( + self, tmp_path: Path, monkeypatch) -> None: + from tensorrt_llm.llmapi._load_balance_env import \ + configure_moe_launch_queues + + config_path = tmp_path / "eplb.yaml" + config_path.write_text("mode: per_layer_dynamic\nnum_slots: 72\n") + monkeypatch.delenv("CUDA_SCALE_LAUNCH_QUEUES", raising=False) + monkeypatch.setattr(torch.cuda, "is_initialized", lambda: False) + + config = MoeConfig(load_balancer=str(config_path)) + overrides = configure_moe_launch_queues(config, {"KEEP": "value"}) + + assert overrides == { + "KEEP": "value", + "CUDA_SCALE_LAUNCH_QUEUES": "4x", + } + assert config.load_balancer.mode == "per_layer_dynamic" + monkeypatch.setattr(torch.cuda, "is_initialized", lambda: True) + assert configure_moe_launch_queues(config, overrides) == overrides + + def test_standard_yaml_does_not_configure_launch_queues( + self, tmp_path: Path, monkeypatch) -> None: + from tensorrt_llm.llmapi._load_balance_env import \ + configure_moe_launch_queues + + config_path = tmp_path / "eplb.yaml" + config_path.write_text("num_slots: 64\nlayer_updates_per_iter: 2\n") + monkeypatch.delenv("CUDA_SCALE_LAUNCH_QUEUES", raising=False) + monkeypatch.setattr(torch.cuda, "is_initialized", lambda: True) + + config = MoeConfig(load_balancer=str(config_path)) + original = {"KEEP": "value"} + + assert configure_moe_launch_queues(config, original) is original + assert "CUDA_SCALE_LAUNCH_QUEUES" not in original + + def test_yaml_resolution_preserves_missing_file_error( + self, tmp_path: Path) -> None: + config_path = tmp_path / "missing.yaml" + + with pytest.raises(FileNotFoundError, + match=r"MoE load balancer config file not found:"): + MoeConfig(load_balancer=str(config_path)) + + def test_yaml_resolution_preserves_invalid_file_error( + self, tmp_path: Path) -> None: + config_path = tmp_path / "invalid.yaml" + config_path.write_text("mode: per_layer_dynamic\n") + + with pytest.raises( + ValidationError, + match=r"Failed to load MoE load balancer config file:"): + MoeConfig(load_balancer=str(config_path)) + + @pytest.mark.cpu_only def test_rejection_sampling_allows_attention_dp(monkeypatch): """ADP (incl. ADP+LM-head-TP) supports rejection sampling. @@ -4127,7 +4328,6 @@ class TestPydanticBestPractices: "drafter", # abstract base class type "resource_manager", # abstract base class type ], - MoeConfig: ["load_balancer"], # allows multiple types including dict RayPlacementConfig: ["placement_groups"], # contains Ray-specific types } diff --git a/tests/unittest/llmapi/test_serialization.py b/tests/unittest/llmapi/test_serialization.py index cbf59b545d19..f8a637fab6eb 100644 --- a/tests/unittest/llmapi/test_serialization.py +++ b/tests/unittest/llmapi/test_serialization.py @@ -24,6 +24,33 @@ def test_serialization_allowed_class(): assert type(obj) == type(b) and obj.name == b.name +@pytest.mark.parametrize( + "config", + [ + pytest.param( + __import__("tensorrt_llm.llmapi.llm_args", + fromlist=["MoeLoadBalancerConfig" + ]).MoeLoadBalancerConfig(num_slots=8), + id="standard", + ), + pytest.param( + __import__( + "tensorrt_llm.llmapi.llm_args", + fromlist=["MoePerLayerDynamicLoadBalancerConfig"], + ).MoePerLayerDynamicLoadBalancerConfig(num_slots=16), + id="per-layer-dynamic", + ), + ], +) +def test_moe_load_balancer_configs_use_restricted_serialization(config): + restored = serialization.loads( + serialization.dumps(config), + approved_imports=serialization.BASE_EXAMPLE_CLASSES, + ) + assert type(restored) is type(config) + assert restored == config + + def test_serialization_disallowed_class(): obj = TestClass("test") a = serialization.dumps(obj) diff --git a/tests/unittest/usage/test_llmapi_config_capture.py b/tests/unittest/usage/test_llmapi_config_capture.py index 13af3416689e..d9dbc29f3043 100644 --- a/tests/unittest/usage/test_llmapi_config_capture.py +++ b/tests/unittest/usage/test_llmapi_config_capture.py @@ -884,15 +884,18 @@ class _RawExcludeConfig(StrictBaseModel): assert config == {"kept": 1} -def test_manifest_excludes_loosely_typed_model_children(): - # B-1 regression: moe_config.load_balancer is Optional[Union[object, str]]; - # a validator coerces it into a MoeLoadBalancerConfig at runtime, but the - # annotation names no BaseModel, so its children must NOT be capturable. +def test_manifest_includes_safe_load_balancer_fields(): from tensorrt_llm.llmapi.llm_args import TorchLlmArgs from tensorrt_llm.usage.llmapi_config import build_capture_manifest - paths = {e.path for e in build_capture_manifest(TorchLlmArgs)} - assert not any(p.startswith("moe_config.load_balancer.") for p in paths) + paths = {entry.path for entry in build_capture_manifest(TorchLlmArgs)} + assert { + "moe_config.load_balancer.auxiliary_sms", + "moe_config.load_balancer.layer_updates_per_iter", + "moe_config.load_balancer.mode", + "moe_config.load_balancer.num_slots", + } <= paths + assert "moe_config.load_balancer.initial_global_assignments" not in paths def test_collect_llm_api_config_caps_total_payload_size(monkeypatch): From 36b96c8b2aa3360806791a390370503b60c7778c Mon Sep 17 00:00:00 2001 From: jintaop Date: Thu, 15 Oct 2026 22:45:10 -0700 Subject: [PATCH 02/59] [None][fix] Address per-layer dynamic EPLB review feedback - Load-balancer seams: model code passes the shared-expert overlap through MoE.load_balancer_overlap(), which _load_balancer_route consumes once; layer balancers declare produces_replica_plan; the no-op set_warmup seam is removed; MoeLoadBalancer and PerLayerDynamicMoeLoadBalancer share BaseMoeLoadBalancer instead of one subclassing the other; the layer handle allocates the live-weight arena and builds its scheduler group from public backend accessors. - MegaMoE: remove the provider-region and dtype-view memo, so the OFF path matches main. - Config: mark the new per-layer dynamic fields as prototype, explain that the extra num_slots are shared helper banks, and describe every load-balancer field (TestPydanticBestPractices). - Docs: document per_layer_dynamic in the wide-EP READMEs and in parallel-strategy.md; rename the remaining user-visible strings to "per-layer dynamic EPLB". - Vendoring: remove the unused VENDOR_STAMP files. - Tests: replace history and source-text pins with behavior checks, stop leaking CUDA_SCALE_LAUNCH_QUEUES, and fix l0_cpu failures after the rebase (rpc_server on fake MPI workers, stale ep_process_group signature test). - Shutdown: remove the unused ep_process_group fallback. Signed-off-by: jintaop --- .../loadBalance/dynamicEplb/CMakeLists.txt | 6 +- .../dynamicEplb/moeRebalancePlaceholder.cpp | 4 +- docs/source/features/parallel-strategy.md | 10 +- examples/wide_ep/README.md | 19 +- examples/wide_ep/ep_load_balancer/README.md | 41 +++ .../_torch/models/modeling_deepseekv4.py | 26 +- .../moe/fused_moe/MOE_DEVELOPER_GUIDE.md | 14 + .../_torch/moe/fused_moe/interface.py | 52 ++- .../fused_moe/mega_moe/mega_moe_cute_dsl.py | 98 ++---- .../mega_moe/rebalance_slot_scheduler.py | 62 ++-- .../_torch/moe/fused_moe/moe_load_balancer.py | 170 +++++---- .../moe/fused_moe/per_layer_dynamic_eplb.py | 126 ++++--- .../moe/kernels/cutedsl_megamoe/VENDOR_STAMP | 7 - .../kernels/megamoe_scheduler/VENDOR_STAMP | 6 - .../_torch/pyexecutor/model_engine.py | 9 +- .../_torch/pyexecutor/model_loader.py | 6 +- tensorrt_llm/llmapi/llm_args.py | 36 +- tensorrt_llm/llmapi/rlhf_utils.py | 4 +- .../_torch/moe/test_dynamic_eplb_contracts.py | 322 ++++++++++-------- tests/unittest/_torch/moe/test_moe_backend.py | 31 +- .../_torch/moe/test_per_layer_dynamic_eplb.py | 127 ++++++- .../executor/test_shutdown_retry_contracts.py | 3 + tests/unittest/llmapi/test_llm_args.py | 6 +- 23 files changed, 738 insertions(+), 447 deletions(-) delete mode 100644 tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/VENDOR_STAMP delete mode 100644 tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/VENDOR_STAMP diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/CMakeLists.txt b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/CMakeLists.txt index 1edbb9c4c973..d2ada7394c9c 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/CMakeLists.txt @@ -30,7 +30,7 @@ else() list(APPEND DYNAMIC_EPLB_DEFINITIONS TRTLLM_DYNAMIC_EPLB_PLACEHOLDER_HALO_Q) message( STATUS - "Dynamic EPLB HALO-Q disabled: none of SM100, SM103, or SM107 was requested" + "Per-layer dynamic EPLB HALO-Q disabled: none of SM100, SM103, or SM107 was requested" ) endif() @@ -43,12 +43,12 @@ else() if(CMAKE_CUDA_COMPILER_VERSION VERSION_LESS 13.1) message( STATUS - "Dynamic EPLB TMA disabled: CUDA 13.1 or newer is required (found ${CMAKE_CUDA_COMPILER_VERSION})" + "Per-layer dynamic EPLB TMA disabled: CUDA 13.1 or newer is required (found ${CMAKE_CUDA_COMPILER_VERSION})" ) else() message( STATUS - "Dynamic EPLB TMA disabled: none of SM100, SM103, or SM107 was requested" + "Per-layer dynamic EPLB TMA disabled: none of SM100, SM103, or SM107 was requested" ) endif() endif() diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlaceholder.cpp b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlaceholder.cpp index e5afc3c3e2d4..dd1b5f075985 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlaceholder.cpp +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlaceholder.cpp @@ -30,7 +30,7 @@ namespace kernels void invokeMoeRebalanceHaloQ(MoeRebalanceHaloQParams const&, cudaStream_t) { - TLLM_THROW("Dynamic EPLB HALO-Q requires a requested SM100, SM103, or SM107 target"); + TLLM_THROW("Per-layer dynamic EPLB HALO-Q requires a requested SM100, SM103, or SM107 target"); } } // namespace kernels @@ -93,7 +93,7 @@ extern "C" char const* megamoe_tma_copy_error_string(int error) { if (error == static_cast(cudaErrorNotSupported)) { - return "Dynamic EPLB TMA requires CUDA 13.1+ and a requested SM100, SM103, or SM107 target"; + return "Per-layer dynamic EPLB TMA requires CUDA 13.1+ and a requested SM100, SM103, or SM107 target"; } return cudaGetErrorString(static_cast(error)); } diff --git a/docs/source/features/parallel-strategy.md b/docs/source/features/parallel-strategy.md index 58732079936c..8a39725f21b6 100644 --- a/docs/source/features/parallel-strategy.md +++ b/docs/source/features/parallel-strategy.md @@ -141,7 +141,7 @@ Large-scale MoE models like DeepSeek-V3/R1, LLaMA4, and Qwen3 use fine-grained e Wide-EP introduces the concept of **expert slots** that are decoupled from specific experts. This allows: - Multiple replicas of hot experts across different GPUs - Dynamic expert placement based on workload patterns -- Both offline and online load balancing strategies +- Offline, online, and per-layer dynamic load balancing strategies #### 2. Custom EP Communication Kernels - Optimized for NVIDIA GB200 Multi-Node NVLink (MNNVL) @@ -151,6 +151,7 @@ Wide-EP introduces the concept of **expert slots** that are decoupled from speci #### 3. Expert Parallelism Load Balancer (EPLB) - **Offline EPLB**: Pre-computed expert placement based on historical workload statistics - **Online EPLB**: Dynamic expert placement that adapts to real-time traffic patterns +- **Per-layer dynamic EPLB**: Expert placement re-planned for every MoE layer on every forward pass, with hot-expert weights copied into shared helper slots on the GPU (`MEGAMOE_CUTEDSL` backend only) - Layer-wise weight redistribution to minimize inference disruption ### Architecture Overview @@ -166,9 +167,10 @@ The system maintains a routing table that maps Expert IDs to Slot IDs, which can 1. **Start with offline EPLB** for production deployments with known workload patterns 2. **Use online EPLB** for dynamic workloads or when traffic patterns change frequently -3. **Monitor expert statistics** to understand workload distribution -4. **Tune max_num_tokens** based on your memory constraints and EP size -5. **Test with representative datasets** to validate load balancing effectiveness +3. **Use per-layer dynamic EPLB** for context-heavy (prefill) serving on the `MEGAMOE_CUTEDSL` backend when the hot experts change from one batch to the next. Online EPLB refreshes a few layers per iteration from decayed statistics, so it follows the recent average load; per-layer dynamic EPLB re-plans each layer from the routing of the current forward pass. It runs without CUDA graphs, so prefer online EPLB for generation-heavy deployments that rely on them. See [Per-Layer Dynamic EP Load Balancer](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/wide_ep/ep_load_balancer#per-layer-dynamic-ep-load-balancer) for its requirements. +4. **Monitor expert statistics** to understand workload distribution +5. **Tune max_num_tokens** based on your memory constraints and EP size +6. **Test with representative datasets** to validate load balancing effectiveness ### References diff --git a/examples/wide_ep/README.md b/examples/wide_ep/README.md index 4cac840c86d7..5770e8a29681 100644 --- a/examples/wide_ep/README.md +++ b/examples/wide_ep/README.md @@ -67,9 +67,9 @@ If set, at most `max_num_tokens` tokens will be sent to `torch.ops.trtllm.fused_ #### `load_balancer` -Configuration for MoE load balancing, users can directly set `num_slots` and `layer_updates_per_iter` as online EPLB settings, while set path to a YAML file that also includes `initial_global_assignments` for offline EPLB. +Configuration for MoE load balancing, users can directly set `num_slots` and `layer_updates_per_iter` as online EPLB settings, while set path to a YAML file that also includes `initial_global_assignments` for offline EPLB. The [`mode`](#mode) field selects between this standard EPLB and per-layer dynamic EPLB. -EPLB is only set up when all of the following hold. Otherwise the `load_balancer` section is silently ignored: +EPLB is only set up when all of the following hold. Otherwise a `standard` `load_balancer` section is silently ignored, while `mode: per_layer_dynamic` raises an error: * The model architecture is listed in `moe_model_arch_list` in [`moe_load_balancer.py`](../../tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py) — DeepSeek-V3/V3.2/V4, GLM MoE DSA, GPT-OSS, Kimi K2.5, Mixtral, Llama 4, Nemotron-H, Qwen2/Qwen3/Qwen3.5 MoE, and the Qwen3-VL / Qwen3.5 MoE conditional-generation wrappers. * Expert parallelism is enabled (`moe_expert_parallel_size` greater than 1). @@ -77,6 +77,13 @@ EPLB is only set up when all of the following hold. Otherwise the `load_balancer Attention DP (`enable_attention_dp: true`) is not part of this gate, so EPLB is also set up without it. It is nonetheless the configuration wide-EP deployments normally use, and it is enabled throughout the [EP Load Balancer examples](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/wide_ep/ep_load_balancer). +#### `mode` + +Load-balancing mode, `standard` or `per_layer_dynamic`. Configs without `mode` are treated as `standard`. + +* `standard`: offline or online EPLB, selected by [`layer_updates_per_iter`](#layer_updates_per_iter). +* `per_layer_dynamic`: per-layer dynamic EPLB, which re-plans the expert replicas of every MoE layer on every forward pass from that pass's routing. It requires the `MEGAMOE_CUTEDSL` backend (`AUTO` resolves to it) and accepts only `num_slots` and `auxiliary_sms` besides `mode`. See [Per-Layer Dynamic EP Load Balancer](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/wide_ep/ep_load_balancer#per-layer-dynamic-ep-load-balancer) for an example and its constraints. + #### `num_slots` Total number of expert slots. It must be ≥ total experts, and divisible by the EP size so that every rank holds the same number of slots. Two typical settings: @@ -85,12 +92,16 @@ Total number of expert slots. It must be ≥ total experts, and divisible by the 2. Set to number of total experts + EP size, such as 288 for DeepSeek R1, 32-way EP. * This means there is 1 extra expert on each EP rank, so that there is more room for the per-rank token distribution to be more balanced. +With `mode: per_layer_dynamic`, `num_slots` is required and must be greater than the total expert count, which must itself be divisible by the EP size. Each rank keeps its `num_experts / ep_size` experts and gets `(num_slots - num_experts) / ep_size` extra slots. Unlike `standard` mode, the extra slots are shared helper banks rather than per-layer replicas: all MoE layers alternate between the same two helper banks, and each layer refills its bank with copies of its hot experts on every forward pass. For example, 384 experts with `num_slots: 416` on 8 EP ranks give each rank 48 experts and 4 helper slots. + To disable MoE load balancing, omit the whole `load_balancer` section. EPLB is enabled by the presence of `load_balancer`, so setting `num_slots` to 0 does not disable it. #### `layer_updates_per_iter` Number of layers updated per iteration, defaults to `0`. `0` means offline, while `>0` means online EPLB. +`layer_updates_per_iter` and `initial_global_assignments` apply only to `standard` mode; a `per_layer_dynamic` configuration that sets either of them is rejected. + This is the number of layers refreshed per iteration, not how often a given layer is refreshed. Layers take turns over `ceil(num_moe_layers / layer_updates_per_iter)` iterations before the plan repeats, so with `layer_updates_per_iter: 1` a DeepSeek-R1 layer (58 MoE layers) is refreshed once every 59 iterations. See [Update Strategy](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/wide_ep/ep_load_balancer#update-strategy) for the statistics, the replication and placement rules, and a worked example. ### Offline Load Balancer Configuration @@ -99,6 +110,10 @@ Refer to the [Offline EP Load Balancer](https://github.com/NVIDIA/TensorRT-LLM/t *Note: Online EP Load Balancer is more suitable for production deployments that need to react timely to online traffic changes.* +### Per-Layer Dynamic Load Balancer Configuration + +Refer to the [Per-Layer Dynamic EP Load Balancer](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/wide_ep/ep_load_balancer#per-layer-dynamic-ep-load-balancer) documentation. + ### Execute Wide-EP on SLURM Clusters Refer to the [slurm_scripts](./slurm_scripts/) directory, which reuses [disaggregated slurm scripts](../disaggregated/slurm/) for submitting jobs to SLURM clusters. diff --git a/examples/wide_ep/ep_load_balancer/README.md b/examples/wide_ep/ep_load_balancer/README.md index e7365d41739e..96e2987fab29 100644 --- a/examples/wide_ep/ep_load_balancer/README.md +++ b/examples/wide_ep/ep_load_balancer/README.md @@ -270,3 +270,44 @@ R3 = 26.7 + 26.7 + 10 = 63.4 ``` The hottest rank is now at 70.0 against a mean of 65, so the imbalance drops from 177% to 7.7%. The three replicas of expert 0 land on three different ranks and share its traffic. This example traces the algorithm on the given input rather than reporting measured performance; for measured results, see the [tech blog](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/blogs/tech_blog/blog04_Scaling_Expert_Parallelism_in_TensorRT-LLM.md#online-ep-load-balancer-1). + + +## Per-Layer Dynamic EP Load Balancer + +Offline EPLB keeps one expert placement for the whole run, and online EPLB refreshes `layer_updates_per_iter` layers per iteration from decayed statistics. Per-layer dynamic EPLB instead re-plans every MoE layer on every forward pass, from the routing of that pass: + +1. Each rank keeps its own `num_experts / ep_size` experts and has `(num_slots - num_experts) / ep_size` helper slots. The helper slots live in two banks that all MoE layers share, rather than in per-layer replicas. +2. Before a layer's experts run, a GPU scheduler (HALO-Q) exchanges routed-token counts across the EP ranks, picks the hot experts to copy into helper slots, and maps every routed token to a resident or helper slot. +3. A copy kernel moves the picked expert weights over NVLink on a high-priority auxiliary stream, while the main stream runs other work such as input quantization and, for DeepSeek-V4, the shared experts. +4. The MoE kernel reads a helper slot only after its weights are marked ready. + +Expert weights are not staged in host shared memory, so the GDRCopy and shared-memory setup that online EPLB needs does not apply. + +An example configuration for a model with 384 routed experts on 8 GPUs: + +```bash +cat > ./config_per_layer_dynamic.yaml < bool: @@ -699,26 +702,45 @@ def _load_balancer_update_statistic(self, is_last_stage=is_last_call, allreduce=self.allreduce) + @contextlib.contextmanager + def load_balancer_overlap(self, + overlap_fn: Callable[[], None]) -> Iterator[None]: + """Run ``overlap_fn`` inside the next load-balancer route call. + + A per-layer dynamic EPLB layer calls it after the copy stream has been + submitted and before the routes are returned, so independent + main-stream work such as shared experts overlaps planning and the + helper-weight copy. Only the first route call consumes it, and it is + cleared on every exit; if no route call consumes it, the caller still + owns that work. Every EP rank must install the same function. + """ + if not callable(overlap_fn): + raise TypeError("overlap_fn must be callable") + self._load_balancer_overlap_fn = overlap_fn + try: + yield + finally: + self._load_balancer_overlap_fn = None + def _load_balancer_route(self, token_selected_experts: torch.Tensor, use_dp: bool) -> torch.Tensor: """Route tokens and enqueue a per-layer dynamic replica plan when enabled.""" if not self.layer_load_balancer: return token_selected_experts - gap_hook = getattr(self, "_rebalance_plan_gap_hook", None) - if gap_hook is not None: - # The hook is one-shot: DeepSeekV4 falls back to ordinary shared - # expert execution if routing exits before consuming it. - self._rebalance_plan_gap_hook = None - route = self.layer_load_balancer.route - if gap_hook is not None: - return route(token_selected_experts, use_dp, gap_hook=gap_hook) - return route(token_selected_experts, use_dp) + overlap_fn = self._load_balancer_overlap_fn + if overlap_fn is None: + return self.layer_load_balancer.route(token_selected_experts, + use_dp) + self._load_balancer_overlap_fn = None + return self.layer_load_balancer.route(token_selected_experts, + use_dp, + overlap_fn=overlap_fn) def _load_balancer_uses_replica_plan(self) -> bool: """Whether routing is produced by the per-layer dynamic device planner.""" - return bool(self.layer_load_balancer and callable( - getattr(self.layer_load_balancer, "require_replica_plan", None))) + return bool(self.layer_load_balancer + and self.layer_load_balancer.produces_replica_plan) def _load_balancer_replica_plan(self): """Return the plan produced by the most recent route, if this mode has one.""" @@ -730,11 +752,7 @@ def _load_balancer_finish_replica_plan(self, plan) -> None: """Release a caller-owned plan after its last MAIN-stream consumer.""" if plan is None or not self.layer_load_balancer: return - finish = getattr(self.layer_load_balancer, "finish_replica_plan", None) - if not callable(finish): - raise RuntimeError( - "an EPLB replica plan was produced without a finish method") - finish(plan) + self.layer_load_balancer.finish_replica_plan(plan) def _load_balancer_start_set_cpu_stage(self, is_last_call: bool): """Start CPU stage in load balancer.""" diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py index 77f19f920bef..04fb4ef0723c 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py @@ -279,7 +279,6 @@ def is_megamoe_cute_dsl_runtime_available() -> Tuple[bool, Optional[str]]: # ``_to_cute``. -# The launch path memoizes these dtype views without copying their storage. def _as_nvfp4(t: torch.Tensor) -> torch.Tensor: return t if t.dtype == torch.float4_e2m1fn_x2 else t.view(torch.float4_e2m1fn_x2) @@ -652,7 +651,6 @@ def __init__( self._per_layer_dynamic_eplb_config = ( load_balancer if getattr(load_balancer, "mode", None) == "per_layer_dynamic" else None ) - self._rebalance_autotune_helpers_initialized = False self._rebalance_slots_active = 0 if self._per_layer_dynamic_eplb_config is not None: ep_size = int(self.ep_size) @@ -684,7 +682,7 @@ def __init__( self._quantize_reserved_sms = self._rebalance_reserved_sms else: logger.warning_once( - "fp4_quantize.sm_budget is unavailable; DynamicEPLB " + "fp4_quantize.sm_budget is unavailable; per-layer dynamic EPLB " "quantization will use the selected backend without SM control", key="dynamic_eplb_quantize_without_sm_budget", ) @@ -774,10 +772,6 @@ def __init__( # scalar would reset the wrong buffer when buckets alternate). self._local_staging_cache: Dict[Tuple, Dict[str, torch.Tensor]] = {} self._last_staged_T: Dict[int, int] = {} - # Cache immutable provider regions and dtype views, retaining their owners. - # Hits validate object identity; dtype views also detect storage rebinding. - self._regions_memo: Dict[int, Tuple] = {} - self._dtype_view_memo: Dict[Tuple, Tuple] = {} if not model_config.skip_create_weights_in_init: self.create_weights() @@ -801,6 +795,24 @@ def compute_slot_count(self) -> int: ) return resident_experts + int(self._rebalance_slots_active) + # ------------------------------------------------------------------ + # Facts the per-layer dynamic EPLB layer handle binds against + # ------------------------------------------------------------------ + @property + def live_weight_arena(self): + """The H + S arena that the live kernel planes alias, or ``None``.""" + return self._rebalance_arena + + @property + def ep_process_group(self): + """The EP ProcessGroup resolved at construction, or ``None`` on one rank.""" + return self._ep_pg + + @property + def max_tokens_per_rank_capacity(self) -> int: + """The largest per-rank token count of any launch bucket.""" + return int(self._maxt_buckets[-1]) + # ------------------------------------------------------------------ # Topology # ------------------------------------------------------------------ @@ -1105,26 +1117,21 @@ def create_weights(self): self._alloc_symm_provider() # Allocate the H + S arena before the quant method creates the seven # kernel-plane Parameters. Source/checkpoint Parameters remain H rows. - if self._rebalance_slots_active > 0 and getattr(self, "_rebalance_arena", None) is None: - # The hierarchical multicast arena: one multicast group per HALO - # level, rather than one per (source, destination) pair. - from .rebalance_live_arena import allocate_rebalance_arena - + if self._rebalance_slots_active > 0 and self._rebalance_arena is None: + # The layer handle allocates the hierarchical multicast arena (one + # multicast group per HALO level, rather than one per source and + # destination pair) from model-owned resources. layer_balancer = getattr(self, "layer_load_balancer", None) - manager = getattr(layer_balancer, "manager", None) - create_provider = getattr(manager, "create_arena_provider", None) - if not callable(create_provider): - raise RuntimeError("per-layer dynamic EPLB requires a model-owned arena provider") - get_ep_comm = getattr(manager, "get_ep_comm", None) - if not callable(get_ep_comm): - raise RuntimeError("per-layer dynamic EPLB requires a model-owned EP communicator") - self._rebalance_arena = allocate_rebalance_arena( - home_experts=int(self._rebalance_home_experts), - helper_slots=int(self._rebalance_slots_active), + allocate_arena = getattr(layer_balancer, "allocate_arena", None) + if not callable(allocate_arena): + raise RuntimeError( + "per-layer dynamic EPLB requires a layer load balancer that " + "allocates the live-weight arena" + ) + self._rebalance_arena = allocate_arena( bundle=live_weight_plane_spec, - ep_comm=get_ep_comm(self._ep_pg), + ep_process_group=self._ep_pg, mapping=self.mapping, - provider=create_provider(), device=int(self.mapping.local_rank), layer_idx=self.layer_idx, ) @@ -1133,7 +1140,7 @@ def create_weights(self): if self._rebalance_arena is not None: from .rebalance_slot_scheduler import bind_live_weight_planes - bind_live_weight_planes(self.quant_method, self) + bind_live_weight_planes(self.quant_method, self, self._rebalance_arena) # Step 4. self._weights_created = True @@ -1630,33 +1637,6 @@ def _ensure_local_staging(self, *, top_k: int, hidden: int, device, output_dtype cached[cache_key] = staging return staging - def _memo_regions(self, provider): - """Cache views of an immutable provider, retaining its identity and lifetime.""" - key = id(provider) - hit = self._regions_memo.get(key) - # Retain the provider and check identity to prevent object-ID reuse. - if hit is not None and hit[0] is provider: - return hit[1] - regions = provider.get_regions() - self._regions_memo[key] = (provider, regions) - return regions - - def _memo_dtype_view(self, t: torch.Tensor, caster) -> torch.Tensor: - """Cache dtype views while retaining the source tensor. - - In-place helper-weight copies remain visible through these aliases. - Storage rebinding must invalidate the cached view even if identity matches. - """ - key = (id(t), caster) - hit = self._dtype_view_memo.get(key) - # Check identity and data_ptr: assigning Parameter.data can replace storage - # without changing the object. These dtype views preserve shape and stride. - if hit is not None and hit[0] is t and hit[1].data_ptr() == t.data_ptr(): - return hit[1] - v = caster(t) - self._dtype_view_memo[key] = (t, v) - return v - def _acquire_buffers( self, *, top_k: int, hidden: int, device, output_dtype, launch_max_T: int ) -> _MegaMoeBuffers: @@ -1708,7 +1688,7 @@ def _acquire_buffers( f"{provider.num_topk} but run_moe called with " f"top_k={top_k}; recreate the backend." ) - regions = self._memo_regions(provider) + regions = provider.get_regions() return _MegaMoeBuffers( activation=regions.activation, activation_sf=regions.activation_sf, @@ -1843,8 +1823,6 @@ def _deferred_scratch_factory(_self=self, _get=get_megamoe_profiling_scratch): if in_kernel_fc2_reduce and num_tokens > 0: combine_output[:num_tokens].zero_() - # Do not memoize scratch-provider regions: a strong reference would keep - # its multi-GiB allocation alive after profiling scratch is released. set_active_megamoe_profiling_scratch( scratch_provider.get_regions() if (scratch_provider is not None and world_size > 1) @@ -1875,16 +1853,14 @@ def _deferred_scratch_factory(_self=self, _get=get_megamoe_profiling_scratch): self._wait_rebalance_routes(replica_plan) try: torch.ops.trtllm.cute_dsl_megamoe_nvfp4_blackwell( - # Activations are per-forward allocations; only persistent weights - # belong in the strong-reference dtype-view memo. activation=_as_nvfp4(activation), activation_sf=_as_fp8_sf(activation_sf), topk_idx=topk_idx, topk_weights=topk_weights, - fc1_weight=self._memo_dtype_view(weight_view.fc1_weight, _as_nvfp4), - fc1_weight_sf=self._memo_dtype_view(weight_view.fc1_weight_sf, _as_fp8_sf), - fc2_weight=self._memo_dtype_view(weight_view.fc2_weight, _as_nvfp4), - fc2_weight_sf=self._memo_dtype_view(weight_view.fc2_weight_sf, _as_fp8_sf), + fc1_weight=_as_nvfp4(weight_view.fc1_weight), + fc1_weight_sf=_as_fp8_sf(weight_view.fc1_weight_sf), + fc2_weight=_as_nvfp4(weight_view.fc2_weight), + fc2_weight_sf=_as_fp8_sf(weight_view.fc2_weight_sf), fc1_alpha=weight_view.fc31_alpha, fc2_alpha=weight_view.fc2_alpha, fc1_norm_const=weight_view.fc1_norm_const, diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py index 2ffa6745cb55..fee74b16533d 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py @@ -419,47 +419,41 @@ def finish(self, completion_event: torch.cuda.Event | None = None) -> None: self._finished_generation = self._generation -def build_rebalance_slot_scheduler_group(backend: Any) -> RebalanceSlotSchedulerGroup: - """Build the layer's scheduler/copy group after its live arena exists.""" - arena = getattr(backend, "_rebalance_arena", None) - if arena is None: - raise RuntimeError( - "MoE rebalance requires create_weights to allocate the live arena " - "before building the producer group." - ) - home = getattr(backend, "_rebalance_home_experts", None) - if type(home) is not int or home <= 0: - raise RuntimeError( - f"MoE rebalance requires a positive resident expert count; got {home!r}." - ) - layer_balancer = getattr(backend, "layer_load_balancer", None) - manager = getattr(layer_balancer, "manager", None) - get_copy_stream = getattr(manager, "get_copy_stream", None) - get_ep_comm = getattr(manager, "get_ep_comm", None) - if not callable(get_copy_stream) or not callable(get_ep_comm): - raise RuntimeError( - "per-layer dynamic EPLB scheduler requires model-owned stream and communicator" - ) +def build_rebalance_slot_scheduler_group( + *, + arena: Any, + manager: Any, + mapping: Any, + ep_process_group: Any, + home_experts: int, + helper_slots: int, + topk: int, + max_tokens_per_rank: int, + layer_idx: Optional[int], +) -> RebalanceSlotSchedulerGroup: + """Build one layer's scheduler/copy group from its manager's resources. + + The manager supplies the model-owned EP communicator, copy stream and + auxiliary SM budget; the communicator is resolved before the stream. + """ + device = int(mapping.local_rank) return RebalanceSlotSchedulerGroup( arena=arena, - mapping=backend.mapping, - ep_comm=get_ep_comm(backend._ep_pg), - home_experts=home, - helper_slots=int(backend._rebalance_slots_active), - max_tokens_per_rank=int(backend._maxt_buckets[-1]), - topk=int(backend.routing_method.experts_per_token), - device=int(backend.mapping.local_rank), - copy_stream=get_copy_stream(int(backend.mapping.local_rank)), + mapping=mapping, + ep_comm=manager.get_ep_comm(ep_process_group), + home_experts=int(home_experts), + helper_slots=int(helper_slots), + max_tokens_per_rank=int(max_tokens_per_rank), + topk=int(topk), + device=device, + copy_stream=manager.get_copy_stream(device), auxiliary_sms=int(manager.auxiliary_sms), - layer_idx=getattr(backend, "layer_idx", None), + layer_idx=layer_idx, ) -def bind_live_weight_planes(quant_method: Any, module: Any) -> None: +def bind_live_weight_planes(quant_method: Any, module: Any, arena: Any) -> None: """Alias MegaMoE's seven live parameters onto its rebalance arena.""" - arena = getattr(module, "_rebalance_arena", None) - if arena is None: - return expected_bundle = quant_method.live_weight_plane_spec(module) if expected_bundle != arena.bundle: raise RuntimeError("MoE rebalance arena does not match the quant-method plane spec") diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py b/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py index 65113c0fa609..69c642d64978 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py @@ -4,6 +4,7 @@ import gc import os import threading +from abc import ABC, abstractmethod from contextlib import AbstractContextManager, nullcontext from multiprocessing import resource_tracker, shared_memory from typing import Callable, Dict, List, Optional, Tuple @@ -303,6 +304,9 @@ class SingleLayerMoeLoadBalancer: This class wraps the C++ implementation of SingleLayerMoeLoadBalancer. """ + # Standard EPLB routes to slot IDs and has no per-invocation replica plan. + produces_replica_plan = False + def __init__( self, single_layer_load_balancer_impl: _tbr.SingleLayerMoeLoadBalancer, @@ -761,11 +765,112 @@ def py_post_shutdown_cleanup(self): self.host_tensor_sharer.post_shutdown_cleanup() -# Global variable to store the current active MoeLoadBalancer instance +# Global variable to store the current active MoE load balancer instance _current_moe_load_balancer = threading.local() -class MoeLoadBalancer: +class BaseMoeLoadBalancer(ABC): + """Model-level lifecycle shared by every MoE load-balancer mode. + + Model construction enters the balancer as a context manager so that MoE + layers register through ``get_moe_load_balancer``. The engine drives each + forward through ``MoeLoadBalancerIterContext`` (``set_iter_info``, + ``start_iter`` and ``end_iter``), and workers release resources with a + collective ``shutdown``. Subclasses initialize ``_previous_balancer``, + ``enable_statistic`` and ``enable_update_weights``. + """ + + @abstractmethod + def is_static_routing(self) -> bool: + """Whether the placement stays fixed, so no iteration hooks run.""" + + def is_dynamic_routing(self) -> bool: + return not self.is_static_routing() + + @abstractmethod + def set_repeated_for_next_layer(self, repeated_count: int) -> None: + """Set how many times the next added layer is invoked per forward.""" + + @abstractmethod + def add_layer(self, + expert_count: int, + top_k: int, + slot_count_per_rank: int, + aux_stream: Optional[torch.cuda.Stream] = None): + """Register one MoE layer and return its layer load balancer.""" + + @abstractmethod + def register_weight_slots_after_to_cuda(self) -> None: + """Register weight slots after the model has been moved to CUDA.""" + + @abstractmethod + def finalize_model(self) -> None: + """Finish model-level setup after every layer has been added.""" + + @abstractmethod + def set_warm_up_iter_count(self, iter_count: int) -> None: + """Set the number of warm-up iterations.""" + + @abstractmethod + def start_iter(self) -> None: + """Start one forward iteration.""" + + @abstractmethod + def end_iter(self) -> None: + """End the current forward iteration.""" + + @abstractmethod + def shutdown(self, local_safe: bool = True) -> bool: + """Collectively release resources; return False if any rank is unsafe.""" + + def set_iter_info(self, enable_statistic: Optional[bool], + enable_update_weights: Optional[bool]): + if enable_statistic is not None: + self.enable_statistic = enable_statistic + if enable_update_weights is not None: + self.enable_update_weights = enable_update_weights + + def __enter__(self): + """ + Enter the context manager. + + Returns: + The load balancer instance + """ + # Save the previous load balancer (if any) + try: + self._previous_balancer = getattr(_current_moe_load_balancer, + 'instance', None) + except AttributeError: + self._previous_balancer = None + + # Set the current instance as the active load balancer + _current_moe_load_balancer.instance = self + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """ + Exit the context manager. + + Args: + exc_type: The exception type + exc_val: The exception value + exc_tb: The exception traceback + + Returns: + False to not suppress exceptions + """ + # Restore the previous load balancer + if self._previous_balancer is not None: + _current_moe_load_balancer.instance = self._previous_balancer + else: + delattr(_current_moe_load_balancer, 'instance') + + # Do not suppress exceptions + return False + + +class MoeLoadBalancer(BaseMoeLoadBalancer): """ A class representing a Mixture of Experts (MoE) load balancer. This class can be used as a context manager to manage the lifecycle of a MoeLoadBalancer. @@ -818,9 +923,6 @@ def is_static_routing(self): # if we don't update, then it is statistic routing. return self.layer_updates_per_iter == 0 - def is_dynamic_routing(self): - return not self.is_static_routing() - def _setup_mpi_comm(self): global_mpi_comm = tensorrt_llm.mpi_comm() shared_mpi_comm = global_mpi_comm.Split_type( @@ -912,13 +1014,6 @@ def set_warm_up_iter_count(self, iter_count: int): """ self.load_balancer_impl.set_warm_up_iter_count(iter_count) - def set_iter_info(self, enable_statistic: Optional[bool], - enable_update_weights: Optional[bool]): - if enable_statistic is not None: - self.enable_statistic = enable_statistic - if enable_update_weights is not None: - self.enable_update_weights = enable_update_weights - def reconfigure_mask_only(self, dead_ranks: list[int]) -> None: """ Reconfigure EPLB routing so slots on dead EP ranks are unreachable. @@ -980,45 +1075,6 @@ def __repr__(self): """ return f"MoeLoadBalancer(ep_rank={self.ep_rank}, ep_size={self.ep_size}, layer_updates_per_iter={self.layer_updates_per_iter})" - def __enter__(self): - """ - Enter the context manager. - - Returns: - The MoeLoadBalancer instance - """ - # Save the previous load balancer (if any) - try: - self._previous_balancer = getattr(_current_moe_load_balancer, - 'instance', None) - except AttributeError: - self._previous_balancer = None - - # Set the current instance as the active load balancer - _current_moe_load_balancer.instance = self - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - """ - Exit the context manager. - - Args: - exc_type: The exception type - exc_val: The exception value - exc_tb: The exception traceback - - Returns: - False to not suppress exceptions - """ - # Restore the previous load balancer - if self._previous_balancer is not None: - _current_moe_load_balancer.instance = self._previous_balancer - else: - delattr(_current_moe_load_balancer, 'instance') - - # Do not suppress exceptions - return False - moe_model_arch_list = [ 'DeepseekV3ForCausalLM', @@ -1049,14 +1105,14 @@ def __exit__(self, exc_type, exc_val, exc_tb): def maybe_create_moe_load_balancer( model_config, mapping: Optional[Mapping] -) -> AbstractContextManager[MoeLoadBalancer | None]: +) -> AbstractContextManager[BaseMoeLoadBalancer | None]: ep_rank = model_config.mapping.moe_ep_rank ep_size = model_config.mapping.moe_ep_size model_arch = model_config.pretrained_config.architectures[0] using_ep = mapping and mapping.moe_ep_size > 1 in_supported_model_arch = model_arch in moe_model_arch_list using_smart_router = mapping and mapping.moe_cluster_size > 1 - moe_load_balancer: AbstractContextManager[MoeLoadBalancer + moe_load_balancer: AbstractContextManager[BaseMoeLoadBalancer | None] = nullcontext() config = model_config.moe_load_balancer mode = getattr(config, "mode", "standard") @@ -1114,7 +1170,7 @@ def maybe_create_moe_load_balancer( class MoeLoadBalancerIterContext: def __init__(self, - moe_load_balancer: Optional[MoeLoadBalancer], + moe_load_balancer: Optional[BaseMoeLoadBalancer], enable_statistic: Optional[bool] = None, enable_updates: Optional[bool] = None): self.moe_load_balancer = moe_load_balancer @@ -1153,12 +1209,12 @@ def __exit__(self, exc_type, exc_val, exc_tb): return False -def get_moe_load_balancer() -> Optional[MoeLoadBalancer]: +def get_moe_load_balancer() -> Optional[BaseMoeLoadBalancer]: """ - Get the current active MoeLoadBalancer instance. + Get the current active MoE load balancer instance. Returns: - The current active MoeLoadBalancer instance, or None if not in a MoeLoadBalancer context + The current active load balancer instance, or None if not in a load-balancer context """ try: return getattr(_current_moe_load_balancer, 'instance', None) diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py index ce8b53558a80..29d8be77b92d 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py @@ -16,12 +16,15 @@ import torch -from .moe_load_balancer import MoeLoadBalancer +from .moe_load_balancer import BaseMoeLoadBalancer class PerLayerDynamicMoeLayerLoadBalancer: """Adapter between the common MoE routing seam and one replica planner.""" + # ``route`` publishes a ``MoEReplicaPlan`` that the caller must finish. + produces_replica_plan = True + def __init__( self, manager: "PerLayerDynamicMoeLoadBalancer", @@ -159,9 +162,38 @@ def bind_runtime(self, *, arena: Any, scheduler_group: Any) -> None: self._scheduler_group = scheduler_group self.manager._register_runtime(self) + def allocate_arena( + self, + *, + bundle: Any, + ep_process_group: Any, + mapping: Any, + device: int, + layer_idx: Optional[int], + ) -> Any: + """Allocate this layer's H + S live-weight arena from model-owned resources. + + Collective across the EP group: every rank calls it from the same point + of the backend's ``create_weights``, after the symmetric-buffer + rendezvous and before the kernel-plane Parameters are registered. + ``bundle`` is the quant method's live-weight plane spec. + """ + from .mega_moe.rebalance_live_arena import allocate_rebalance_arena + + return allocate_rebalance_arena( + home_experts=self.resident_experts_per_rank, + helper_slots=self.helper_slots_per_rank, + bundle=bundle, + ep_comm=self.manager.get_ep_comm(ep_process_group), + mapping=mapping, + provider=self.manager.create_arena_provider(), + device=int(device), + layer_idx=layer_idx, + ) + def bind_backend(self, backend: Any) -> None: """Build and own the producer after the backend's live aliases exist.""" - arena = getattr(backend, "_rebalance_arena", None) + arena = backend.live_weight_arena if arena is None: raise RuntimeError("per-layer dynamic EPLB backend has no live arena") if self._arena is not None or self._scheduler_group is not None: @@ -196,9 +228,20 @@ def bind_backend(self, backend: Any) -> None: for plane in arena.tekit_alias_views: for helper in range(int(arena.helper_slots)): plane[home + helper].copy_(plane[helper % home], non_blocking=True) - backend._rebalance_autotune_helpers_initialized = True - scheduler_group = build_rebalance_slot_scheduler_group(backend) + # ``layer_idx`` is the backend's model layer index: shared-slot pools + # compare producers across ranks by it. + scheduler_group = build_rebalance_slot_scheduler_group( + arena=arena, + manager=self.manager, + mapping=backend.mapping, + ep_process_group=backend.ep_process_group, + home_experts=self.resident_experts_per_rank, + helper_slots=self.helper_slots_per_rank, + topk=self.top_k, + max_tokens_per_rank=backend.max_tokens_per_rank_capacity, + layer_idx=backend.layer_idx, + ) self.bind_runtime(arena=arena, scheduler_group=scheduler_group) self._bound_backend_ref = backend_ref @@ -211,30 +254,34 @@ def route( token_selected_experts: torch.Tensor, offset_by_ep_rank: bool = False, *, - gap_hook: Optional[Callable[[], None]] = None, + overlap_fn: Optional[Callable[[], None]] = None, ) -> torch.Tensor: """Submit HALO-Q/TMA and return its physical routes without a host wait. - A shared-expert hook runs after the copy stream is submitted but before - the physical routes are lent to the caller, preserving the intended - scheduler/copy overlap. + ``overlap_fn`` runs after the copy stream is submitted but before the + physical routes are lent to the caller, so independent main-stream work + such as shared experts overlaps planning and the helper-weight copy. """ del offset_by_ep_rank # HALO-Q consumes global logical expert IDs. - if gap_hook is not None and not callable(gap_hook): - raise TypeError("gap_hook must be callable") + if overlap_fn is not None and not callable(overlap_fn): + raise TypeError("overlap_fn must be callable") if self._scheduler_group is None or self._arena is None: - raise RuntimeError(f"dynamic EPLB layer {self.layer_idx} has no bound runtime") + raise RuntimeError( + f"per-layer dynamic EPLB layer {self.layer_idx} has no bound runtime" + ) if self._pending_replica_plan is not None: - raise RuntimeError(f"dynamic EPLB layer {self.layer_idx} has an unfinished plan") + raise RuntimeError( + f"per-layer dynamic EPLB layer {self.layer_idx} has an unfinished plan" + ) - if gap_hook is None: + if overlap_fn is None: physical, generation = self._scheduler_group.plan( token_selected_experts, defer_wait=True ) else: part = self._scheduler_group.plan_schedule(token_selected_experts) try: - gap_hook() + overlap_fn() except BaseException as error: # noqa: BLE001 try: self._scheduler_group.discard_plan(part) @@ -248,7 +295,7 @@ def route( if type(notes) is not list: notes = [] error.__notes__ = notes - notes.append("Shared-expert hook failed after HALO-Q/TMA enqueue; " + cleanup_note) + notes.append("overlap_fn failed after HALO-Q/TMA enqueue; " + cleanup_note) raise physical, generation = self._scheduler_group.plan_finish(part, defer_wait=True) generation = int(generation) @@ -269,7 +316,9 @@ def route( def require_replica_plan(self) -> Any: if self._pending_replica_plan is None: - raise RuntimeError(f"dynamic EPLB layer {self.layer_idx} has not produced a plan") + raise RuntimeError( + f"per-layer dynamic EPLB layer {self.layer_idx} has not produced a plan" + ) return self._pending_replica_plan def finish_replica_plan( @@ -280,7 +329,9 @@ def finish_replica_plan( ) -> None: pending = self._pending_replica_plan if pending is None: - raise RuntimeError(f"dynamic EPLB layer {self.layer_idx} has no plan to finish") + raise RuntimeError( + f"per-layer dynamic EPLB layer {self.layer_idx} has no plan to finish" + ) if plan is not None and plan is not pending: raise RuntimeError("attempted to finish a different replica plan") assert self._scheduler_group is not None @@ -308,7 +359,7 @@ def _clear_runtime(self) -> None: self._bound_backend_ref = None -class PerLayerDynamicMoeLoadBalancer(MoeLoadBalancer): +class PerLayerDynamicMoeLoadBalancer(BaseMoeLoadBalancer): """Model-level owner for per-layer dynamic EPLB layers and GPU resources.""" def __init__( @@ -318,10 +369,9 @@ def __init__( ep_size: int, config: Any, mapping: Optional[Any] = None, - ep_process_group: Optional[Any] = None, ) -> None: - # Deliberately do not call MoeLoadBalancer.__init__: per-layer dynamic EPLB - # has no C++ standard-EPLB object and must not create an MPI split. + # Per-layer dynamic EPLB has no C++ standard-EPLB object and creates no + # MPI split; its EP communicator is resolved by the first MoE layer. if ep_size <= 1 or not 0 <= ep_rank < ep_size: raise ValueError(f"invalid EPLB EP rank/size: rank={ep_rank}, size={ep_size}") num_slots = int(config.num_slots) @@ -335,15 +385,13 @@ def __init__( self.ep_size = int(ep_size) self.config = config self.mapping = mapping - self.ep_process_group = ep_process_group self.num_slots = num_slots self.num_local_slots = num_slots // ep_size self.auxiliary_sms = int(getattr(config, "auxiliary_sms", 8)) if self.auxiliary_sms <= 0: raise ValueError("auxiliary_sms must be positive") - # Compatibility state used by existing context/engine call sites. - self.layer_updates_per_iter = 1 + # State read by the shared lifecycle and by the engine. self._previous_balancer = None self.single_layer_load_balancers: list[PerLayerDynamicMoeLayerLoadBalancer] = [] self.next_layer_repeated_count: Optional[int] = None @@ -366,12 +414,16 @@ def __del__(self) -> None: # a Python finalizer cannot safely enter EP collectives. pass + def __repr__(self) -> str: + return ( + f"PerLayerDynamicMoeLoadBalancer(ep_rank={self.ep_rank}, " + f"ep_size={self.ep_size}, num_slots={self.num_slots}, " + f"auxiliary_sms={self.auxiliary_sms})" + ) + def is_static_routing(self) -> bool: return False - def is_dynamic_routing(self) -> bool: - return True - def set_repeated_for_next_layer(self, repeated_count: int) -> None: if repeated_count <= 0: raise ValueError("repeat count must be positive") @@ -452,16 +504,6 @@ def _register_runtime(self, layer: PerLayerDynamicMoeLayerLoadBalancer) -> None: assert layer._arena is not None self.adopt_arena_provider(layer._arena.provider) - def set_iter_info( - self, - enable_statistic: Optional[bool], - enable_update_weights: Optional[bool], - ) -> None: - if enable_statistic is not None: - self.enable_statistic = bool(enable_statistic) - if enable_update_weights is not None: - self.enable_update_weights = bool(enable_update_weights) - def start_iter(self) -> None: if self.in_iter: raise RuntimeError("already in a MoE load-balancer iteration") @@ -480,10 +522,6 @@ def end_iter(self) -> None: self.in_iter = False self.iter_id += 1 - def set_warmup(self, value: bool) -> None: - """Keep production HALO-Q/TMA planning active during warmup.""" - del value - def register_weight_slots_after_to_cuda(self) -> None: # Live aliases are registered by the per-layer dynamic arena binder. pass @@ -517,11 +555,11 @@ def shutdown(self, local_safe: bool = True) -> bool: return True # Agree on local readiness before any rank starts destructive teardown. - # This also covers managers that failed before their first pool existed. + # Every MoE layer resolves the EP communicator in create_weights, and a + # rank whose model load failed never owns an engine that reaches this + # method, so the communicator exists on every EP rank or on none. local_safe = bool(local_safe and not self.in_iter) comm = self._ep_comm - if comm is None and self.ep_process_group is not None: - comm = self.get_ep_comm(self.ep_process_group) if comm is not None: states = comm.allgather( { diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/VENDOR_STAMP b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/VENDOR_STAMP deleted file mode 100644 index 3fbb6a2e0395..000000000000 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/VENDOR_STAMP +++ /dev/null @@ -1,7 +0,0 @@ -upstream: https://gitlab-master.nvidia.com/bangyus/cutedsl_megamoe.git -locked_commit: 382ce4a6f111d0769f31c4c975e1c643f0dc4f70 -project: CuTe DSL MegaMoE -original_authors: Bangyu Shen ; Aragorn Guan ; Hao Sheng ; Yuhan Li ; Yucheng Yan ; XingFei Xi ; Shawn Wang -license: per-file SPDX declarations (BSD-3-Clause where declared) -downstream_patch: 3rdparty/vendor_patches/cutedsl-megamoe.patch (TensorRT-LLM main adaptations: SwiGLU alpha/beta for MiniMax-M3, bandit B105 annotations) -managed_by: 3rdparty/vendor_sources.lock.yaml diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/VENDOR_STAMP b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/VENDOR_STAMP deleted file mode 100644 index f8e67c1f1393..000000000000 --- a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/VENDOR_STAMP +++ /dev/null @@ -1,6 +0,0 @@ -upstream: https://gitlab-master.nvidia.com/jintaop/cutedsl_eplb_scheduler_copy.git -locked_commit: 4dc8616031b18dd948bdf11b49afe0d86be15e22 -project: MegaMoE HALO-Q scheduler and SAMI runtime -original_authors: Jintao Peng ; Ethan Zhang -license: Apache-2.0 -managed_by: 3rdparty/vendor_sources.lock.yaml diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 60d6be2ba9de..2b90bcd9d917 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -42,7 +42,7 @@ from ..models.modeling_utils import DecoderModelForCausalLM, timing_metric from ..moe.expert_statistic import ExpertStatistic from ..moe.fused_moe.moe_comm_timeout_guard import set_moe_comm_warmup -from ..moe.fused_moe.moe_load_balancer import MoeLoadBalancer +from ..moe.fused_moe.moe_load_balancer import BaseMoeLoadBalancer from ..nccl_window_tensor_scope import install_eager_nccl_window_tensor_scopes from ..route_capture import ROUTE_CAPTURE_ATTR, RouteCapture from ..speculative import SpecMetadata, update_spec_config_from_loaded_model @@ -270,7 +270,7 @@ def _load_model( model_weights_restore_mode, ) -> None: """Load the model, or take the given one, at the requested capacity.""" - self.moe_load_balancer: MoeLoadBalancer | None = None + self.moe_load_balancer: BaseMoeLoadBalancer | None = None if model is None: lora_config: LoraConfig | None = self.llm_args.lora_config # Keep the model_loader to support reloading the model weights later @@ -289,7 +289,7 @@ def _load_model( self.model, moe_load_balancer = self.model_loader.load( checkpoint_dir=model_path, checkpoint_loader=checkpoint_loader ) - if isinstance(moe_load_balancer, MoeLoadBalancer): + if isinstance(moe_load_balancer, BaseMoeLoadBalancer): self.moe_load_balancer = moe_load_balancer else: self.model = model @@ -936,9 +936,6 @@ def _set_warmup_effects(self, in_warmup: bool) -> None: self.moe_load_balancer.set_iter_info( enable_statistic=not in_warmup, enable_update_weights=not in_warmup ) - set_warmup = getattr(self.moe_load_balancer, "set_warmup", None) - if callable(set_warmup): - set_warmup(in_warmup) @contextmanager def warmup_scope(self) -> Iterator[None]: diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 94735ae2d0bb..f0061a4b20f0 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -56,7 +56,7 @@ get_registered_model_class, timing_metric) from ..modules.embedding import Embedding, LMHead from ..modules.low_m_gemm import LOW_M_GEMM_ACTIVE, prepare_low_m_gemm -from ..moe.fused_moe.moe_load_balancer import (MoeLoadBalancer, +from ..moe.fused_moe.moe_load_balancer import (BaseMoeLoadBalancer, maybe_create_moe_load_balancer) from ..virtual_memory import RestoreMode from ..virtual_memory import scope as virtual_memory_scope @@ -1070,7 +1070,7 @@ def load( self, checkpoint_dir: str, checkpoint_loader: BaseCheckpointLoader, - ) -> tuple[DecoderModelForCausalLM, MoeLoadBalancer | None]: + ) -> tuple[DecoderModelForCausalLM, BaseMoeLoadBalancer | None]: """ Loads the model, its weights, and applies necessary configurations. @@ -1595,7 +1595,7 @@ def init_meta_tensor_in_pool(t: torch.Tensor): # `tensorrt_llm/llmapi/llm_args.py`. When implementing the # follow-up, drop the validator gate AFTER moving these calls # into the pool scope. - if isinstance(moe_load_balancer, MoeLoadBalancer): + if isinstance(moe_load_balancer, BaseMoeLoadBalancer): moe_load_balancer.register_weight_slots_after_to_cuda() logger.info("moe_load_balancer finalizing model...") moe_load_balancer.finalize_model() diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 69d8dcd9c982..d490c407fac0 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1853,7 +1853,11 @@ def to_sparse_params(self, **kwargs): class _MoeLoadBalancerConfigBase(StrictBaseModel): """Common slot layout and expert-parallel runtime state for EPLB.""" - num_slots: Optional[int] = None + num_slots: Optional[int] = Field( + default=None, + description= + "Total number of expert slots across all EP ranks. Must be divisible by `ep_size`; each rank hosts `num_slots / ep_size` slots." + ) _ep_rank: Optional[int] = PrivateAttr(default=None) _ep_size: Optional[int] = PrivateAttr(default=None) @@ -1903,12 +1907,22 @@ def slot_end(self) -> int: class MoeLoadBalancerConfig(_MoeLoadBalancerConfigBase): """Configure the standard asynchronous EPLB implementation.""" - mode: Literal["standard"] = "standard" + mode: Literal["standard"] = Field( + default="standard", + description= + "EPLB mode. `standard` is static EPLB when `layer_updates_per_iter` is 0 and online EPLB otherwise. Configs without `mode` use `standard`.", + status="prototype") initial_global_assignments: Optional[Dict[int, List[int]]] = Field( default=None, repr=False, + description= + "Initial expert placement keyed by MoE layer index. Each value has `num_slots` entries, gives the expert ID held by each global slot, and covers every expert." + ) + layer_updates_per_iter: int = Field( + default=0, + description= + "Number of MoE layers whose expert placement is updated per iteration. 0 keeps the initial placement (static EPLB); a positive value enables online EPLB." ) - layer_updates_per_iter: int = 0 def get_layer_initial_global_assignments( self, layer_idx: int) -> Optional[List[int]]: @@ -1932,14 +1946,22 @@ def get_layer_initial_global_assignments( class MoePerLayerDynamicLoadBalancerConfig(_MoeLoadBalancerConfigBase): """Configure per-layer dynamic EPLB with Halo-Q scheduling.""" - mode: Literal["per_layer_dynamic"] = "per_layer_dynamic" - num_slots: PositiveInt + mode: Literal["per_layer_dynamic"] = Field( + default="per_layer_dynamic", + description= + "EPLB mode. `per_layer_dynamic` replicates each MoE layer's hot experts into helper slots in every forward pass, planned from that layer's current routing.", + status="prototype") + num_slots: PositiveInt = Field( + description= + "Total number of expert slots across all EP ranks. Must be divisible by `ep_size` and greater than the number of experts. Each rank keeps its `num_experts / ep_size` resident experts and adds `(num_slots - num_experts) / ep_size` helper slots. Unlike standard EPLB, the helper slots are not per-layer replica weights: they are two shared helper banks that all MoE layers alternate between and refill in every forward pass.", + status="prototype") auxiliary_sms: int = Field( default=8, ge=8, le=128, - description="SM budget shared serially by HALO-Q and in-switch copy.", - ) + description= + "Number of SMs that the routed-input quantization and shared-expert FC12 kernels leave free for the HALO-Q scheduler and the in-switch copy kernel, which run one after the other on the copy stream. Must be smaller than the device SM count.", + status="prototype") def validate_expert_capacity(self, num_experts: int) -> None: """Validate that every EP rank has at least one helper slot.""" diff --git a/tensorrt_llm/llmapi/rlhf_utils.py b/tensorrt_llm/llmapi/rlhf_utils.py index 9a9188a18cd8..139b0af15be4 100644 --- a/tensorrt_llm/llmapi/rlhf_utils.py +++ b/tensorrt_llm/llmapi/rlhf_utils.py @@ -8,7 +8,7 @@ import torch -from tensorrt_llm._torch.moe.fused_moe.moe_load_balancer import MoeLoadBalancer +from tensorrt_llm._torch.moe.fused_moe.moe_load_balancer import BaseMoeLoadBalancer from tensorrt_llm._torch.utils import get_device_uuid from tensorrt_llm.executor.ray.utils import control_action_decorator from tensorrt_llm.llmapi import serialization @@ -211,7 +211,7 @@ def update_weights(self, ipc_handles: Optional[dict] = None): logger.info("Finalize update weights") self.finalize_weight_update() moe_load_balancer = getattr(self.engine.model_engine, "moe_load_balancer", None) - if isinstance(moe_load_balancer, MoeLoadBalancer): + if isinstance(moe_load_balancer, BaseMoeLoadBalancer): moe_load_balancer.register_weight_slots_after_to_cuda() logger.info("moe_load_balancer finalizing model...") moe_load_balancer.finalize_model() diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index 6ecf315192ba..49b1bbc98f22 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Fast CPU contracts for the Dynamic EPLB serving path.""" +"""Fast CPU contracts for the per-layer dynamic EPLB serving path.""" from __future__ import annotations @@ -25,7 +25,13 @@ _QUEUE = "CUDA_SCALE_LAUNCH_QUEUES" -def _extract_class(path: Path, name: str, methods: set[str] | None, namespace: dict): +def _extract_class( + path: Path, + name: str, + methods: set[str] | None, + namespace: dict, + bases: Tuple[str, ...] = (), +): tree = ast.parse(path.read_text()) original = next( node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == name @@ -39,7 +45,13 @@ def _extract_class(path: Path, name: str, methods: set[str] | None, namespace: d if isinstance(node, ast.FunctionDef) and node.name in methods ] assert {node.name for node in body} == methods - selected = ast.ClassDef(name=name, bases=[], keywords=[], body=body, decorator_list=[]) + selected = ast.ClassDef( + name=name, + bases=[ast.Name(id=base, ctx=ast.Load()) for base in bases], + keywords=[], + body=body, + decorator_list=[], + ) module = ast.Module( body=[ ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), @@ -51,6 +63,23 @@ def _extract_class(path: Path, name: str, methods: set[str] | None, namespace: d return namespace[name] +def _extract_function(path: Path, name: str, namespace: dict): + tree = ast.parse(path.read_text()) + function = next( + node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == name + ) + function.decorator_list = [] + module = ast.Module( + body=[ + ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), + function, + ], + type_ignores=[], + ) + exec(compile(ast.fix_missing_locations(module), str(path), "exec"), namespace) + return namespace[name] + + def test_shared_slot_provider_uses_model_owned_pool_state(): source = _TORCH_ROOT / "moe/fused_moe/mega_moe/rebalance_live_arena.py" provider_cls = _extract_class( @@ -427,7 +456,6 @@ def test_model_engine_warmup_uses_load_balancer_public_api(): enable_statistic=True, enable_update_weights=True, set_iter_info=lambda **kwargs: events.append(("iter", kwargs)), - set_warmup=lambda value: events.append(("warmup", value)), ) engine = engine_cls() engine.moe_load_balancer = manager @@ -437,7 +465,6 @@ def test_model_engine_warmup_uses_load_balancer_public_api(): assert events == [ ("comm", True), ("iter", {"enable_statistic": False, "enable_update_weights": False}), - ("warmup", True), ] @@ -879,159 +906,154 @@ def config(mode="per_layer_dynamic"): assert os.environ.get(_QUEUE) == process_queue -def test_launch_queue_setup_precedes_runtime_initialization(): - def call_name(node): - if isinstance(node, ast.Name): - return node.id - if isinstance(node, ast.Attribute): - return f"{call_name(node.value)}.{node.attr}" - if ( - isinstance(node, ast.Call) - and isinstance(node.func, ast.Name) - and node.func.id == "super" - ): - return "super()" - return "" - - def calls(relative, function_name, class_name=None): - tree = ast.parse((_ROOT / relative).read_text()) - owner = tree - if class_name is not None: - owner = next( - node - for node in tree.body - if isinstance(node, ast.ClassDef) and node.name == class_name - ) - function = next( - node - for node in owner.body - if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) - and node.name == function_name - ) - result = {} - for node in ast.walk(function): - if isinstance(node, ast.Call): - name = call_name(node.func) - result[name] = min(result.get(name, node.lineno), node.lineno) - return result - - base = calls("tensorrt_llm/llmapi/llm.py", "__init__", "BaseLLM") - assert ( - base["self._process_env_overrides"] - < base["configure_moe_launch_queues"] - < base["llm_args_cls"] - ) - assert base["configure_moe_launch_queues"] < base["get_device_count"] +class _RuntimeInitReached(Exception): + """Stops an entrypoint at its first runtime-initialization call.""" + + +class _RecordingEnviron(dict): + def __init__(self, events, **values): + super().__init__(values) + self._events = events - worker_main = calls("tensorrt_llm/executor/worker.py", "worker_main") - assert ( - worker_main["os.environ.update"] - < worker_main["configure_moe_launch_queues"] - < worker_main["worker_cls"] + def __setitem__(self, key, value): + self._events.append(("env", key, value)) + super().__setitem__(key, value) + + def update(self, values): + for key, value in values.items(): + self[key] = value + + +def test_launch_queue_setup_sees_env_overrides_before_runtime_initialization(monkeypatch): + events = [] + environ = _RecordingEnviron(events, MASTER_ADDR="127.0.0.1") + fake_os = SimpleNamespace( + environ=environ, + getenv=lambda key, default=None: environ.get(key, default), ) - assert worker_main["configure_moe_launch_queues"] < worker_main["mpi_comm"] + logger = SimpleNamespace( + level="warning", + set_level=lambda level: None, + info=lambda *args, **kwargs: None, + error=lambda *args, **kwargs: None, + debug=lambda *args, **kwargs: None, + ) + user_overrides = {"USER_SETTING": "1"} + configured = {**user_overrides, _QUEUE: "4x"} + + def configure_moe_launch_queues(moe_config, env_overrides): + # The override must already be visible when the launch queue is configured. + events.append(("configure", environ.get("USER_SETTING"))) + return {**(env_overrides or {}), _QUEUE: "4x"} + + def runtime_init(name): + def call(*args, **kwargs): + events.append(name) + raise _RuntimeInitReached(name) + + return call + + def run(entrypoint, stop): + events.clear() + environ.pop("USER_SETTING", None) + with pytest.raises(_RuntimeInitReached, match=stop): + entrypoint() + return list(events) + + def worker_args(): + return SimpleNamespace( + backend="pytorch", env_overrides=dict(user_overrides), moe_config=object() + ) + + expected = [("env", "USER_SETTING", "1"), ("configure", "1")] + common = { + "os": fake_os, + "logger": logger, + "configure_moe_launch_queues": configure_moe_launch_queues, + } - worker = calls("tensorrt_llm/executor/base_worker.py", "__init__", "BaseWorker") - assert ( - worker["os.environ.update"] - < worker["configure_moe_launch_queues"] - < worker["super().__init__"] + # LLM API: argument validation is the first step that may probe CUDA. + class TorchLlmArgs: + model_fields = {"env_overrides": None, "moe_config": None} + + def __init__(self, **kwargs): + events.append(("llm_args", kwargs["env_overrides"])) + raise _RuntimeInitReached("llm_args") + + llm_cls = _extract_class( + _ROOT / "tensorrt_llm/llmapi/llm.py", + "BaseLLM", + {"__init__", "_process_env_overrides"}, + { + **common, + "GenerationExecutor": object, + "MoeConfig": lambda **kwargs: SimpleNamespace(**kwargs), + "TorchLlmArgs": TorchLlmArgs, + "TORCH_LLMARGS_REMOVED_ARGS": set(), + "mpi_disabled": lambda: False, + }, ) + assert run( + lambda: llm_cls( + "model", backend="pytorch", env_overrides=dict(user_overrides), moe_config={} + ), + "llm_args", + ) == expected + [("llm_args", configured)] - ray = calls("tensorrt_llm/executor/ray/gpu_worker.py", "__init__", "RayWorkerWrapper") - assert ( - ray["os.environ.update"] - < ray["configure_moe_launch_queues"] - < ray["torch.cuda.device_count"] - < ray["torch.cuda.set_device"] + # MPI worker process: MPI is the first runtime initialization. + worker_main = _extract_function( + _ROOT / "tensorrt_llm/executor/worker.py", + "worker_main", + {**common, "GenerationExecutorWorker": object, "mpi_comm": runtime_init("mpi_comm")}, ) + args = worker_args() + assert run( + lambda: worker_main(engine="engine", worker_queues=None, log_level="info", llm_args=args), + "mpi_comm", + ) == expected + ["mpi_comm"] + assert args.env_overrides == configured + # Direct worker construction: the executor base initializes the runtime. + class GenerationExecutor: + __init__ = runtime_init("executor_init") -def test_routed_megamoe_direct_inputs_and_full_grid_contract(): - source = _TORCH_ROOT / "moe/fused_moe/mega_moe/mega_moe_cute_dsl.py" - tree = ast.parse(source.read_text()) - launch = next( - node - for node in ast.walk(tree) - if isinstance(node, ast.Call) - and isinstance(node.func, ast.Attribute) - and node.func.attr == "cute_dsl_megamoe_nvfp4_blackwell" + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/base_worker.py", + "BaseWorker", + {"__init__"}, + { + **common, + "GenerationExecutor": GenerationExecutor, + "PostprocWorkerConfig": lambda: SimpleNamespace( + num_postprocess_workers=0, postprocess_tokenizer_dir=None + ), + }, + bases=("GenerationExecutor",), ) - arguments = {keyword.arg: keyword.value for keyword in launch.keywords} - - for name, caster in (("activation", "_as_nvfp4"), ("activation_sf", "_as_fp8_sf")): - value = arguments[name] - assert isinstance(value, ast.Call) - assert isinstance(value.func, ast.Name) and value.func.id == caster - - for name in ("fc1_weight", "fc1_weight_sf", "fc2_weight", "fc2_weight_sf"): - value = arguments[name] - assert isinstance(value, ast.Call) - assert isinstance(value.func, ast.Attribute) - assert value.func.attr == "_memo_dtype_view" - - assert "reserved_sms" not in arguments - - -def test_gpu_direct_tma_fails_closed_on_stale_scheduler_plan(): - source = ( - _ROOT / "cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu" - ).read_text() - begin = source.index("__global__ void tma_copy_gpu_direct_kernel") - end = source.index("struct MegamoeTmaCopyState", begin) - kernel = source[begin:end] - - status_check = kernel.index("statusOk &= direct.status[field] == 0") - epoch_check = kernel.index("schedulerPlanReady = statusOk && planEpoch == schedulerEpoch") - rejection = kernel.index("if (!schedulerPlanReady)") - trap = kernel.index('asm volatile("trap;"', rejection) - publish = kernel.index("publish_plan_channel_warp") - build = kernel.index("megamoe_tma_build_gpu_plan") - ready = kernel.index("finish_cta_and_notify") - - assert status_check < epoch_check < rejection <= trap < publish < build < ready - - submit = source[source.index("megamoe_tma_copy_submit_gpu_direct") :] - assert "tma_copy_gpu_direct_kernel<<<" in submit - assert "cudaLaunchAttributeProgrammaticStreamSerialization" not in submit - assert "cudaLaunchKernelEx" not in submit - - -def test_halo_q_rendezvous_resets_arrivals_and_wraps_epoch_safely(): - source = ( - _ROOT / "cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.cu" - ).read_text() - begin = source.index("__device__ bool grid_rendezvous") - end = source.index("__device__ void local_histogram_pass", begin) - rendezvous = source[begin:end] - - assert "using Epoch = std::uint32_t;" in source - assert "return observed != expected && expected - observed < (Epoch{1} << 31);" in source - arrival = rendezvous.index("fetch_add_acq_rel_device(sync, 1)") - last_cta = rendezvous.index("old == static_cast(blocks - 1)") - reset = rendezvous.index("store_relaxed_device(sync, 0)") - publish = rendezvous.index("store_release_device(sync + 1, target)") - wait = rendezvous.index("while (epoch_before(observed, target))") - assert arrival < last_cta < reset < publish < wait - - mask = (1 << 32) - 1 - - def epoch_before(observed: int, expected: int) -> bool: - delta = (expected - observed) & mask - return observed != expected and delta < (1 << 31) - - epochs = (0xFFFFFFFE, 0xFFFFFFFF, 0x00000000, 0x00000001) - arrivals = 0 - for observed, expected in zip(epochs, epochs[1:]): - for _ in range(128): - old = arrivals - arrivals += 1 - if old == 127: - arrivals = 0 - assert arrivals == 0 - assert epoch_before(observed, expected) - assert not epoch_before(expected, observed) - assert not epoch_before(expected, expected) + args = worker_args() + assert run(lambda: worker_cls(engine="engine", llm_args=args), "executor_init") == expected + [ + "executor_init" + ] + assert args.env_overrides == configured + + # Ray worker: the first CUDA query follows launch-queue setup. + monkeypatch.setitem(sys.modules, "tensorrt_llm.logger", SimpleNamespace(logger=logger)) + ray_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/ray/gpu_worker.py", + "RayWorkerWrapper", + {"__init__"}, + { + **common, + "torch": SimpleNamespace( + cuda=SimpleNamespace(device_count=runtime_init("torch.cuda.device_count")) + ), + }, + ) + args = worker_args() + assert run(lambda: ray_cls(object, {"llm_args": args}, 1, 0), "device_count") == expected + [ + "torch.cuda.device_count" + ] + assert args.env_overrides == configured def test_direct_submit_orders_generations_and_defers_the_route_wait(): @@ -1464,7 +1486,6 @@ def close_collectively(self, *, local_safe): manager.is_shutdown = False manager.in_iter = True manager._ep_comm = Comm() - manager.ep_process_group = object() manager._rebalance_shared_slot_pool_order = [Pool()] assert manager.shutdown(local_safe=True) is False @@ -1583,6 +1604,7 @@ def terminal_cleanup(*, local_safe): ) worker = worker_cls() worker.doing_shutdown = False + worker.rpc_server = None worker.engine = engine worker.llm_args = SimpleNamespace(backend="pytorch") worker._handle_background_error = lambda: events.append("background") diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index e86408bdde48..91e1b359a738 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -17,7 +17,6 @@ import collections import dataclasses import importlib -import inspect import itertools import logging import os @@ -2117,16 +2116,14 @@ def test_megamoe_rejects_expand_before_provider_allocation(monkeypatch): moe.ep_size = 2 allocate_provider = MagicMock() moe._alloc_symm_provider = allocate_provider - create_arena_provider = MagicMock() - moe.layer_load_balancer = SimpleNamespace( - manager=SimpleNamespace(create_arena_provider=create_arena_provider) - ) + allocate_arena = MagicMock() + moe.layer_load_balancer = SimpleNamespace(allocate_arena=allocate_arena) with pytest.raises(ValueError, match="expand_intermediate == 2"): moe.create_weights() allocate_provider.assert_not_called() - create_arena_provider.assert_not_called() + allocate_arena.assert_not_called() @pytest.mark.cpu_only @@ -2336,6 +2333,10 @@ def test_megamoe_cutedsl_cache_derived_state_reuses_rebalance_binding( ) moe._rebalance_slots_active = 2 moe._rebalance_arena = arena + moe._ep_pg = object() + moe._maxt_buckets = [64, 128] + moe.mapping = SimpleNamespace(local_rank=0) + moe.layer_idx = 5 moe.layer_load_balancer = layer builder = MagicMock(return_value=SimpleNamespace(home_experts=4, helper_slots=2)) monkeypatch.setattr( @@ -2350,7 +2351,17 @@ def test_megamoe_cutedsl_cache_derived_state_reuses_rebalance_binding( assert moe.quant_method.cache_derived_state.call_count == 2 assert arena.assert_identity.call_count == 2 assert moe.quant_method.assert_live_weight_aliases.call_count == 2 - builder.assert_called_once_with(moe) + builder.assert_called_once_with( + arena=arena, + manager=manager, + mapping=moe.mapping, + ep_process_group=moe._ep_pg, + home_experts=4, + helper_slots=2, + topk=2, + max_tokens_per_rank=128, + layer_idx=5, + ) @pytest.mark.cpu_only @@ -2415,12 +2426,6 @@ def record_event(event_handle, stream_handle): assert trace[-1] == "prove_reuse" -@pytest.mark.cpu_only -def test_rebalance_scheduler_group_requires_ep_process_group(): - parameters = inspect.signature(RebalanceSlotSchedulerGroup.__init__).parameters - assert "ep_process_group" in parameters - - def test_megamoe_load_weights_invalidates_cached_deepgemm_views(): method = W4A8MXFP4MXFP8MegaMoEDeepGemmMethod() hidden_size = 128 diff --git a/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py b/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py index 69fa28889b6f..1259d2e131ea 100644 --- a/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py +++ b/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py @@ -7,10 +7,13 @@ import pytest import torch +from tensorrt_llm._torch.moe.fused_moe.interface import MoE from tensorrt_llm._torch.moe.fused_moe.mega_moe.rebalance_slot_scheduler import ( RebalanceSlotSchedulerGroup, ) from tensorrt_llm._torch.moe.fused_moe.moe_load_balancer import ( + BaseMoeLoadBalancer, + MoeLoadBalancer, MoeLoadBalancerIterContext, get_moe_load_balancer, maybe_create_moe_load_balancer, @@ -63,8 +66,12 @@ def assert_identity(self): class _Backend: def __init__(self, arena): - self._rebalance_arena = arena + self.live_weight_arena = arena self.tactic_autotune = False + self.mapping = SimpleNamespace(local_rank=0) + self.ep_process_group = object() + self.max_tokens_per_rank_capacity = 256 + self.layer_idx = 9 class _Group: @@ -156,7 +163,7 @@ def test_layer_produces_and_finishes_replica_plan(): assert layer.replica_plan is None -def test_gap_hook_runs_between_schedule_and_plan_finish(): +def test_overlap_fn_runs_between_schedule_and_plan_finish(): manager = _manager() layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) arena = _Arena() @@ -164,13 +171,13 @@ def test_gap_hook_runs_between_schedule_and_plan_finish(): layer.bind_runtime(arena=arena, scheduler_group=group) logical = torch.tensor([[0, 1]], dtype=torch.int32) - physical = layer.route(logical, gap_hook=lambda: group.events.append("hook")) + physical = layer.route(logical, overlap_fn=lambda: group.events.append("overlap")) torch.testing.assert_close(physical, logical + 4) - assert group.events == ["schedule", "hook", "finish"] + assert group.events == ["schedule", "overlap", "finish"] layer.finish_replica_plan() -def test_gap_hook_failure_discards_generation_and_adds_context(): +def test_overlap_fn_failure_discards_generation_and_adds_context(): manager = _manager() layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) group = _Group() @@ -180,12 +187,45 @@ def fail(): raise RuntimeError("shared expert failed") with pytest.raises(RuntimeError, match="shared expert failed") as caught: - layer.route(torch.tensor([[0, 1]], dtype=torch.int32), gap_hook=fail) + layer.route(torch.tensor([[0, 1]], dtype=torch.int32), overlap_fn=fail) assert group.events == ["schedule", "discard"] assert "safely released" in caught.value.__notes__[0] assert layer.replica_plan is None +def test_load_balancer_overlap_is_consumed_once_and_always_cleared(): + calls = [] + + class _Layer: + produces_replica_plan = True + + def route(self, routes, offset_by_ep_rank, *, overlap_fn=None): + calls.append(overlap_fn) + return routes + + moe = SimpleNamespace(layer_load_balancer=_Layer(), _load_balancer_overlap_fn=None) + routes = torch.tensor([[0, 1]], dtype=torch.int32) + + def overlap(): + pass + + with MoE.load_balancer_overlap(moe, overlap): + MoE._load_balancer_route(moe, routes, False) + MoE._load_balancer_route(moe, routes, False) + assert calls == [overlap, None] + assert moe._load_balancer_overlap_fn is None + + with pytest.raises(RuntimeError, match="routed path failed"): + with MoE.load_balancer_overlap(moe, overlap): + raise RuntimeError("routed path failed") + assert moe._load_balancer_overlap_fn is None + + with pytest.raises(TypeError, match="overlap_fn must be callable"): + with MoE.load_balancer_overlap(moe, None): + pass + assert MoE._load_balancer_uses_replica_plan(moe) + + def test_layer_disables_standard_host_migration_and_checks_geometry(): manager = _manager() layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) @@ -221,7 +261,19 @@ def test_bind_backend_is_idempotent_before_scheduler_construction(monkeypatch): layer.bind_backend(backend) layer.bind_backend(backend) - builder.assert_called_once_with(backend) + # The group is built from the layer's geometry, the manager's resources and + # the backend's public facts; layer_idx stays the backend's model index. + builder.assert_called_once_with( + arena=arena, + manager=manager, + mapping=backend.mapping, + ep_process_group=backend.ep_process_group, + home_experts=4, + helper_slots=2, + topk=2, + max_tokens_per_rank=256, + layer_idx=9, + ) assert layer._arena is arena assert layer._scheduler_group is group @@ -245,11 +297,11 @@ def test_bind_backend_rejects_different_backend_or_arena_before_construction( with pytest.raises(RuntimeError, match="different backend"): layer.bind_backend(_Backend(arena)) - backend._rebalance_arena = _Arena() + backend.live_weight_arena = _Arena() with pytest.raises(RuntimeError, match="changed its rebalance arena"): layer.bind_backend(backend) - builder.assert_called_once_with(backend) + builder.assert_called_once() def test_iteration_context_reuses_common_lifecycle(): @@ -342,6 +394,63 @@ def validate_geometry(self, *, expected_rank, expected_world): manager.get_ep_comm(object()) +def test_layer_allocates_arena_from_model_owned_resources(monkeypatch): + from tensorrt_llm._torch.moe.fused_moe.mega_moe import rebalance_live_arena + + manager = _manager() + layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) + order = [] + process_group, comm, provider, arena = object(), object(), object(), object() + + def get_ep_comm(group): + assert group is process_group + order.append("ep_comm") + return comm + + def create_arena_provider(): + order.append("provider") + return provider + + monkeypatch.setattr(manager, "get_ep_comm", get_ep_comm) + monkeypatch.setattr(manager, "create_arena_provider", create_arena_provider) + allocate = MagicMock(return_value=arena) + monkeypatch.setattr(rebalance_live_arena, "allocate_rebalance_arena", allocate) + bundle, mapping = object(), object() + + allocated = layer.allocate_arena( + bundle=bundle, + ep_process_group=process_group, + mapping=mapping, + device=1, + layer_idx=7, + ) + + assert allocated is arena + assert order == ["ep_comm", "provider"] + allocate.assert_called_once_with( + home_experts=4, + helper_slots=2, + bundle=bundle, + ep_comm=comm, + mapping=mapping, + provider=provider, + device=1, + layer_idx=7, + ) + + +def test_per_layer_dynamic_manager_shares_the_base_lifecycle_only(): + manager = _manager() + assert isinstance(manager, BaseMoeLoadBalancer) + assert not isinstance(manager, MoeLoadBalancer) + assert manager.is_dynamic_routing() + assert repr(manager) == ( + "PerLayerDynamicMoeLoadBalancer(ep_rank=0, ep_size=4, num_slots=24, auxiliary_sms=8)" + ) + layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) + assert layer.produces_replica_plan + + def test_factory_selects_per_layer_dynamic_manager(): mapping = SimpleNamespace( moe_ep_rank=0, diff --git a/tests/unittest/executor/test_shutdown_retry_contracts.py b/tests/unittest/executor/test_shutdown_retry_contracts.py index f446c559a29d..f5485fd8c33d 100644 --- a/tests/unittest/executor/test_shutdown_retry_contracts.py +++ b/tests/unittest/executor/test_shutdown_retry_contracts.py @@ -421,6 +421,7 @@ def release_symm_buffer_cache(): ) worker = worker_cls() worker.doing_shutdown = False + worker.rpc_server = None worker.engine = engine worker.llm_args = SimpleNamespace(backend="pytorch") worker._handle_background_error = lambda: events.append("background") @@ -501,6 +502,7 @@ def allgather(self, value): ) worker = worker_cls() worker.doing_shutdown = False + worker.rpc_server = None worker.engine = engine worker.llm_args = SimpleNamespace(backend="pytorch") worker._handle_background_error = lambda: events.append("background") @@ -563,6 +565,7 @@ def test_mpi_worker_background_error_keeps_shutdown_committed(monkeypatch): ) worker = worker_cls() worker.doing_shutdown = False + worker.rpc_server = None worker.engine = SimpleNamespace( shutdown_all_ranks=True, can_enqueue_requests=lambda: False, diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index bcaf206da232..84046cf7f422 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -401,7 +401,11 @@ def test_per_layer_dynamic_yaml_configures_queue_before_cuda_init( config_path = tmp_path / "eplb.yaml" config_path.write_text("mode: per_layer_dynamic\nnum_slots: 72\n") - monkeypatch.delenv("CUDA_SCALE_LAUNCH_QUEUES", raising=False) + # The code under test writes os.environ. Seed the key first so that + # monkeypatch records the caller's state and removes "4x" on teardown; + # delenv alone records nothing when the key is absent. + monkeypatch.setenv("CUDA_SCALE_LAUNCH_QUEUES", "") + monkeypatch.delenv("CUDA_SCALE_LAUNCH_QUEUES") monkeypatch.setattr(torch.cuda, "is_initialized", lambda: False) config = MoeConfig(load_balancer=str(config_path)) From 9ea179d0dda45e2cd732238754ebab619b8a092e Mon Sep 17 00:00:00 2001 From: jintaop Date: Fri, 16 Oct 2026 00:16:31 -0700 Subject: [PATCH 03/59] [None][fix] Accept per-layer dynamic EPLB in the serve config schemas The serve config schema accepted only standard EPLB mappings for moe_config.load_balancer, so editors flagged valid per_layer_dynamic configs. Accept both forms, and require mode only for per-layer dynamic EPLB because a mapping without mode is standard EPLB. Check the accepted and rejected forms against MoeConfig. Regenerate the checked-in schemas. This also picks up KvCacheConfig.kv_pool_rebalance_config and ParallelConfig.tp_layout, which main added without regenerating the snapshots. Signed-off-by: jintaop --- scripts/generate_trtllm_serve_schemas.py | 16 +- .../schemas/trtllm-serve-config.schema.json | 142 ++++++++++++------ .../trtllm-serve-disagg-config.schema.json | 142 ++++++++++++------ .../test_trtllm_serve_config_schemas.py | 30 ++++ 4 files changed, 242 insertions(+), 88 deletions(-) diff --git a/scripts/generate_trtllm_serve_schemas.py b/scripts/generate_trtllm_serve_schemas.py index a4e3522d8dd1..07bd42e8c8f9 100644 --- a/scripts/generate_trtllm_serve_schemas.py +++ b/scripts/generate_trtllm_serve_schemas.py @@ -123,6 +123,7 @@ def generate_serve_schema() -> dict: from tensorrt_llm.llmapi.llm_args import ( SERVE_CONFIG_ALIASES, MoeLoadBalancerConfig, + MoePerLayerDynamicLoadBalancerConfig, TorchLlmArgs, ) @@ -141,10 +142,17 @@ def generate_serve_schema() -> dict: ) # Serving treats null as unset before constructing the non-optional telemetry model. properties["telemetry_config"] = _nullable(properties["telemetry_config"]) - # Keep accepted forms aligned with TorchLlmArgs.validate_load_balancer(). - schema["$defs"]["MoeConfig"]["properties"]["load_balancer"] = { - "anyOf": [_schema_for(MoeLoadBalancerConfig, schema), {"type": ["string", "null"]}] - } + # Keep accepted forms aligned with TorchLlmArgs.validate_load_balancer(): a mapping + # without mode is standard EPLB, so only the per-layer dynamic form requires it. + standard, per_layer_dynamic = ( + {"$ref": f"#/$defs/{model.__name__}"} + for model in (MoeLoadBalancerConfig, MoePerLayerDynamicLoadBalancerConfig) + ) + schema["$defs"]["MoeConfig"]["properties"]["load_balancer"]["anyOf"] = [ + standard, + {**per_layer_dynamic, "required": ["mode"]}, + {"type": ["string", "null"]}, + ] # These object/Any annotations are Python API hooks, not arbitrary YAML mappings. # Update these restrictions when hooks are added or gain YAML construction support. for name in ("batched_logits_processor", "checkpoint_loader", "_mpi_session"): diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index b0e84576f7bf..5d22437c9b32 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -4032,47 +4032,13 @@ "load_balancer": { "anyOf": [ { - "additionalProperties": false, - "description": "Pydantic configuration model for the Mixture of Experts (MoE) load balancer.\n\nThis model holds configuration data (`num_slots`, etc.) as well as\nruntime state (`_ep_rank`, `_ep_size`) which must be set via the\n`setup()` method before use.", - "properties": { - "initial_global_assignments": { - "anyOf": [ - { - "additionalProperties": { - "items": { - "type": "integer" - }, - "type": "array" - }, - "type": "object" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Initial Global Assignments" - }, - "layer_updates_per_iter": { - "default": 0, - "title": "Layer Updates Per Iter", - "type": "integer" - }, - "num_slots": { - "anyOf": [ - { - "type": "integer" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Num Slots" - } - }, - "title": "MoeLoadBalancerConfig", - "type": "object" + "$ref": "#/$defs/MoeLoadBalancerConfig" + }, + { + "$ref": "#/$defs/MoePerLayerDynamicLoadBalancerConfig", + "required": [ + "mode" + ] }, { "type": [ @@ -4080,7 +4046,10 @@ "null" ] } - ] + ], + "default": null, + "description": "Configuration for standard or per-layer dynamic EPLB.", + "title": "Load Balancer" }, "max_num_tokens": { "anyOf": [ @@ -4105,6 +4074,95 @@ "title": "MoeConfig", "type": "object" }, + "MoeLoadBalancerConfig": { + "additionalProperties": false, + "description": "Configure the standard asynchronous EPLB implementation.", + "properties": { + "initial_global_assignments": { + "anyOf": [ + { + "additionalProperties": { + "items": { + "type": "integer" + }, + "type": "array" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Initial expert placement keyed by MoE layer index. Each value has `num_slots` entries, gives the expert ID held by each global slot, and covers every expert.", + "title": "Initial Global Assignments" + }, + "layer_updates_per_iter": { + "default": 0, + "description": "Number of MoE layers whose expert placement is updated per iteration. 0 keeps the initial placement (static EPLB); a positive value enables online EPLB.", + "title": "Layer Updates Per Iter", + "type": "integer" + }, + "mode": { + "const": "standard", + "default": "standard", + "description": "EPLB mode. `standard` is static EPLB when `layer_updates_per_iter` is 0 and online EPLB otherwise. Configs without `mode` use `standard`.", + "status": "prototype", + "title": "Mode", + "type": "string" + }, + "num_slots": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Total number of expert slots across all EP ranks. Must be divisible by `ep_size`; each rank hosts `num_slots / ep_size` slots.", + "title": "Num Slots" + } + }, + "title": "MoeLoadBalancerConfig", + "type": "object" + }, + "MoePerLayerDynamicLoadBalancerConfig": { + "additionalProperties": false, + "description": "Configure per-layer dynamic EPLB with Halo-Q scheduling.", + "properties": { + "auxiliary_sms": { + "default": 8, + "description": "Number of SMs that the routed-input quantization and shared-expert FC12 kernels leave free for the HALO-Q scheduler and the in-switch copy kernel, which run one after the other on the copy stream. Must be smaller than the device SM count.", + "maximum": 128, + "minimum": 8, + "status": "prototype", + "title": "Auxiliary Sms", + "type": "integer" + }, + "mode": { + "const": "per_layer_dynamic", + "default": "per_layer_dynamic", + "description": "EPLB mode. `per_layer_dynamic` replicates each MoE layer's hot experts into helper slots in every forward pass, planned from that layer's current routing.", + "status": "prototype", + "title": "Mode", + "type": "string" + }, + "num_slots": { + "description": "Total number of expert slots across all EP ranks. Must be divisible by `ep_size` and greater than the number of experts. Each rank keeps its `num_experts / ep_size` resident experts and adds `(num_slots - num_experts) / ep_size` helper slots. Unlike standard EPLB, the helper slots are not per-layer replica weights: they are two shared helper banks that all MoE layers alternate between and refill in every forward pass.", + "exclusiveMinimum": 0, + "status": "prototype", + "title": "Num Slots", + "type": "integer" + } + }, + "required": [ + "num_slots" + ], + "title": "MoePerLayerDynamicLoadBalancerConfig", + "type": "object" + }, "MooncakeStoreConfig": { "additionalProperties": false, "description": "How this server joins a Mooncake store pool.\n\nThe settings every participant shares come from the master named by\n`pool`; what is left here is per-server. Setting this makes\n`trtllm-serve` render the Mooncake client config and export\n`MOONCAKE_CONFIG_PATH`; an inherited `MOONCAKE_CONFIG_PATH` wins.", diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index abca43a18d20..0ce7858017f6 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -5596,47 +5596,13 @@ "load_balancer": { "anyOf": [ { - "additionalProperties": false, - "description": "Pydantic configuration model for the Mixture of Experts (MoE) load balancer.\n\nThis model holds configuration data (`num_slots`, etc.) as well as\nruntime state (`_ep_rank`, `_ep_size`) which must be set via the\n`setup()` method before use.", - "properties": { - "initial_global_assignments": { - "anyOf": [ - { - "additionalProperties": { - "items": { - "type": "integer" - }, - "type": "array" - }, - "type": "object" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Initial Global Assignments" - }, - "layer_updates_per_iter": { - "default": 0, - "title": "Layer Updates Per Iter", - "type": "integer" - }, - "num_slots": { - "anyOf": [ - { - "type": "integer" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Num Slots" - } - }, - "title": "MoeLoadBalancerConfig", - "type": "object" + "$ref": "#/$defs/MoeLoadBalancerConfig" + }, + { + "$ref": "#/$defs/MoePerLayerDynamicLoadBalancerConfig", + "required": [ + "mode" + ] }, { "type": [ @@ -5644,7 +5610,10 @@ "null" ] } - ] + ], + "default": null, + "description": "Configuration for standard or per-layer dynamic EPLB.", + "title": "Load Balancer" }, "max_num_tokens": { "anyOf": [ @@ -5669,6 +5638,95 @@ "title": "MoeConfig", "type": "object" }, + "MoeLoadBalancerConfig": { + "additionalProperties": false, + "description": "Configure the standard asynchronous EPLB implementation.", + "properties": { + "initial_global_assignments": { + "anyOf": [ + { + "additionalProperties": { + "items": { + "type": "integer" + }, + "type": "array" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Initial expert placement keyed by MoE layer index. Each value has `num_slots` entries, gives the expert ID held by each global slot, and covers every expert.", + "title": "Initial Global Assignments" + }, + "layer_updates_per_iter": { + "default": 0, + "description": "Number of MoE layers whose expert placement is updated per iteration. 0 keeps the initial placement (static EPLB); a positive value enables online EPLB.", + "title": "Layer Updates Per Iter", + "type": "integer" + }, + "mode": { + "const": "standard", + "default": "standard", + "description": "EPLB mode. `standard` is static EPLB when `layer_updates_per_iter` is 0 and online EPLB otherwise. Configs without `mode` use `standard`.", + "status": "prototype", + "title": "Mode", + "type": "string" + }, + "num_slots": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Total number of expert slots across all EP ranks. Must be divisible by `ep_size`; each rank hosts `num_slots / ep_size` slots.", + "title": "Num Slots" + } + }, + "title": "MoeLoadBalancerConfig", + "type": "object" + }, + "MoePerLayerDynamicLoadBalancerConfig": { + "additionalProperties": false, + "description": "Configure per-layer dynamic EPLB with Halo-Q scheduling.", + "properties": { + "auxiliary_sms": { + "default": 8, + "description": "Number of SMs that the routed-input quantization and shared-expert FC12 kernels leave free for the HALO-Q scheduler and the in-switch copy kernel, which run one after the other on the copy stream. Must be smaller than the device SM count.", + "maximum": 128, + "minimum": 8, + "status": "prototype", + "title": "Auxiliary Sms", + "type": "integer" + }, + "mode": { + "const": "per_layer_dynamic", + "default": "per_layer_dynamic", + "description": "EPLB mode. `per_layer_dynamic` replicates each MoE layer's hot experts into helper slots in every forward pass, planned from that layer's current routing.", + "status": "prototype", + "title": "Mode", + "type": "string" + }, + "num_slots": { + "description": "Total number of expert slots across all EP ranks. Must be divisible by `ep_size` and greater than the number of experts. Each rank keeps its `num_experts / ep_size` resident experts and adds `(num_slots - num_experts) / ep_size` helper slots. Unlike standard EPLB, the helper slots are not per-layer replica weights: they are two shared helper banks that all MoE layers alternate between and refill in every forward pass.", + "exclusiveMinimum": 0, + "status": "prototype", + "title": "Num Slots", + "type": "integer" + } + }, + "required": [ + "num_slots" + ], + "title": "MoePerLayerDynamicLoadBalancerConfig", + "type": "object" + }, "MooncakeStoreConfig": { "additionalProperties": false, "description": "How this server joins a Mooncake store pool.\n\nThe settings every participant shares come from the master named by\n`pool`; what is left here is per-server. Setting this makes\n`trtllm-serve` render the Mooncake client config and export\n`MOONCAKE_CONFIG_PATH`; an inherited `MOONCAKE_CONFIG_PATH` wins.", diff --git a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py index 592287377cf6..f25058eddd27 100644 --- a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py +++ b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py @@ -60,6 +60,7 @@ def schemas(tmp_path_factory: pytest.TempPathFactory) -> dict[str, dict]: (generator.SERVE_SCHEMA, {"moe_config": {"backend": "CUTLASS"}}, True), (generator.SERVE_SCHEMA, {"moe_config": {"backend_typo": "CUTLASS"}}, False), (generator.SERVE_SCHEMA, {"moe_config": {"backend": "not-a-backend"}}, False), + (generator.SERVE_SCHEMA, {"moe_config": {"load_balancer": "eplb.yaml"}}, True), (generator.SERVE_SCHEMA, {"backend": "pytorch"}, True), (generator.SERVE_SCHEMA, {"backend": "_autodeploy"}, False), ( @@ -203,6 +204,35 @@ def test_branch_snapshot_field_types( assert (not errors) == valid, [error.message for error in errors] +@pytest.mark.parametrize( + ("load_balancer", "valid"), + [ + ({"num_slots": 288}, True), + ({"mode": "standard", "num_slots": 288, "layer_updates_per_iter": 1}, True), + ({"mode": "per_layer_dynamic", "num_slots": 320, "auxiliary_sms": 16}, True), + # A mapping without mode is standard EPLB, which has no auxiliary_sms. + ({"num_slots": 320, "auxiliary_sms": 16}, False), + ({"mode": "per_layer_dynamic"}, False), + ({"mode": "per_layer_dynamic", "num_slots": 320, "layer_updates_per_iter": 1}, False), + ({"mode": "per_layer_dynamic", "num_slots": 320, "auxiliary_sms": 4}, False), + ({"mode": "dynamic", "num_slots": 320}, False), + ], +) +def test_load_balancer_schema_matches_runtime( + schemas: dict[str, dict], load_balancer: dict, valid: bool +) -> None: + from tensorrt_llm.llmapi.llm_args import MoeConfig + + config = {"moe_config": {"load_balancer": load_balancer}} + errors = list(Draft202012Validator(schemas[generator.SERVE_SCHEMA]).iter_errors(config)) + assert (not errors) == valid, [error.message for error in errors] + if valid: + MoeConfig(load_balancer=load_balancer) + else: + with pytest.raises(ValueError): + MoeConfig(load_balancer=load_balancer) + + def _assert_serving_fields_in_schema(source: str, function_name: str, properties: dict) -> None: """Check literal top-level config accesses, without tracing aliases or dynamic keys.""" function = next( From 66604a0fec874a523289c82e51f0cf8a92423e22 Mon Sep 17 00:00:00 2001 From: jintaop Date: Sun, 18 Oct 2026 07:33:00 -0700 Subject: [PATCH 04/59] [None][chore] Update vendored cutedsl_megamoe to bangyus main 7832a330 Move the cutedsl-megamoe vendor from bangyus/cutedsl_megamoe main 382ce4a6 to main 7832a330. The selected file set is unchanged. The lock and compatibility patch were updated only with scripts/vendor_sources.py: the existing patch no longer applies to 7832a330, so the destination was set to the exact upstream selection and pinned with "pin" (which drops a patch that cannot apply when the destination is exact upstream), and the TensorRT-LLM compatibility edits were then re-applied and recorded with "patch create". Upstream commits that touch the vendored files: - 43af6d6c Report deterministic router overflow group-wide (token_comm_deterministic.py; merged via 822e739b) - d921a6b6 router: partition the metadata push by route index (MR !78) (token_comm.py) - 13434224 publish fc1/fc2 done with red.async.release (epi_async_release; Blackwell fc12 epilogue and MegaMoE kernel) - bd19a17b Blackwell bf16 swapab kernel for training forward pass (BF16 wire, token_in_stage_split, external_dispatch_pool in token_comm.py and token_comm_deterministic.py) - ce9173db Fix GenPhase workspace reset ordering (Rubin genphase MegaMoE kernel) - 76676a54 FC1 per-token activation scale for the Blackwell nvfp4 MegaMoE path (opt-in; also adds gate_activation="quick_geglu" and the megamoe_aot trailing optional argument) Compatibility patch, re-applied on 7832a330 (TensorRT-LLM-only): - SwiGLU swiglu_alpha/swiglu_beta support in the Blackwell and Rubin MegaMoE kernels and epilogues (MiniMax-M3). The Blackwell epilogue now derives its sigmoid exp2 constant from a quick_geglu scale; swiglu_alpha sets that scale, giving the same -alpha * log2e constant as before, and gate_activation="quick_geglu" is rejected together with SwiGLUBias constants, as SiTU already is. - "# nosec B105" on token_* region names and mode-string comparisons for bandit. The new upstream code needs no additional annotations. - Removal of "if stream_ends_before_target: fc1_exhausted = Boolean(True)" in PhaseInterleavedFc12Scheduler, as on TensorRT-LLM main (#19960). - Removal of the perf-number comment in helpers/topk_reduce.py. TensorRT-LLM integration needs no adaptation. The new descriptor keys (gate_activation, fc1_activation_per_token_scale, epi_async_release, token_in_stage_split, external_dispatch_pool) and the trailing fc1_activation_per_token_scale kernel argument are optional, and TensorRT-LLM does not set them. Blackwell MegaMoE and the shared FC12 kernel therefore use the upstream default epi_async_release=True, which requires epilogue flag batches <= 4; every TensorRT-LLM tactic uses (1, 1), (1, 4) or (2, 4). Workspace sizes are unchanged for TensorRT-LLM configurations. Signed-off-by: jintaop --- 3rdparty/vendor_patches/cutedsl-megamoe.patch | 173 +++---- 3rdparty/vendor_sources.lock.yaml | 6 +- .../communication/nvlink_domain/token_comm.py | 432 ++++++++++++------ .../nvlink_domain/token_comm_deterministic.py | 338 +++++++++----- .../cutedsl_megamoe/helpers/megamoe_aot.py | 6 +- .../block_scaled_swap_ab_fc12_epilogue.py | 313 ++++++++++++- .../block_scaled_swap_ab_fc12_extension.py | 5 +- .../mega/block_scaled_swap_ab_fc12_kernel.py | 25 +- .../block_scaled_swap_ab_mega_moe_kernel.py | 56 ++- .../block_scaled_swap_ab_fc12_epilogue.py | 5 + .../block_scaled_swap_ab_fc12_epilogue.py | 6 + ...swap_ab_mega_moe_kernel_gen_specialized.py | 62 +-- 12 files changed, 1046 insertions(+), 381 deletions(-) diff --git a/3rdparty/vendor_patches/cutedsl-megamoe.patch b/3rdparty/vendor_patches/cutedsl-megamoe.patch index cfc9eb9edb10..9c594f14ec2c 100644 --- a/3rdparty/vendor_patches/cutedsl-megamoe.patch +++ b/3rdparty/vendor_patches/cutedsl-megamoe.patch @@ -1,17 +1,17 @@ diff --git a/communication/nvlink_domain/token_comm.py b/communication/nvlink_domain/token_comm.py -index 7e2d5b2a70509a1a960fa450f95642ef321c9913..910e753adb0fcff4969ed89903d4d250cab06850 100644 +index a64a5e7e377b77373863a766db4a1ed70c9b9252..374aa01b055bb187dd588c64ec0482a54c9d1d32 100644 --- a/communication/nvlink_domain/token_comm.py +++ b/communication/nvlink_domain/token_comm.py -@@ -113,7 +113,7 @@ class _MetadataPushRouter(KernelComponent): +@@ -123,7 +123,7 @@ class _MetadataPushRouter(KernelComponent): sorted_metadata_region = "nvlink.token_comm.sorted_metadata" sorted_scores_region = "nvlink.token_comm.sorted_scores" pool_expert_base_region = "nvlink.token_comm.pool_expert_base" - token_src_metadata_region = "nvlink.token_comm.token_src_metadata" + token_src_metadata_region = "nvlink.token_comm.token_src_metadata" # nosec B105 fc1_topk_scores_region = "nvlink.token_comm.fc1_topk_scores" - source_expert_base_region = "nvlink.token_comm.source_expert_base" - push_destination_base_region = "nvlink.token_comm.push_destination_base" -@@ -1087,16 +1087,16 @@ class TokenCommNonDeterministic(KernelComponent): + # Pool-ordered fp32 per-token FC1 activation scale, pushed next to the metadata like the scores. + fc1_activation_scales_region = "nvlink.token_comm.fc1_activation_per_token_scales" +@@ -1178,16 +1178,16 @@ class TokenCommNonDeterministic(KernelComponent): fc2_activation_sf_region = "nvlink.token_comm.fc2_activation_sf" pre_reduced_activation_region = "nvlink.token_comm.pre_reduced_activation" pre_reduced_activation_sf_region = "nvlink.token_comm.pre_reduced_activation_sf" @@ -36,7 +36,7 @@ index 7e2d5b2a70509a1a960fa450f95642ef321c9913..910e753adb0fcff4969ed89903d4d250 @classmethod def problem_desc_require(cls) -> dict[str, type]: -@@ -1184,7 +1184,7 @@ class TokenCommNonDeterministic(KernelComponent): +@@ -1277,7 +1277,7 @@ class TokenCommNonDeterministic(KernelComponent): # FC12 occupies warps 0..7. Legacy primitive callers retain their # modulo-four numbering by omitting the explicit CTA ordinal. self.dispatch_warp_start = 8 @@ -45,7 +45,7 @@ index 7e2d5b2a70509a1a960fa450f95642ef321c9913..910e753adb0fcff4969ed89903d4d250 self.dispatch_warp_count = self.transfer_warp_count self.combine_warp_count = self.transfer_warp_count self.dispatch_cta_count = self.promised_launchable_sm_count -@@ -1261,9 +1261,9 @@ class TokenCommNonDeterministic(KernelComponent): +@@ -1356,9 +1356,9 @@ class TokenCommNonDeterministic(KernelComponent): ) if self.token_back_ready_granularity not in ("expert", "token_tile"): raise ValueError(f"Unsupported token_back_ready_granularity {self.token_back_ready_granularity!r}.") @@ -57,8 +57,8 @@ index 7e2d5b2a70509a1a960fa450f95642ef321c9913..910e753adb0fcff4969ed89903d4d250 ): raise ValueError("token_tile return-ready requires standalone_warps or reuse_dispatch_warps with atomic_counter.") if not 1 <= self.token_in_flag_batch <= 32: -@@ -1281,7 +1281,7 @@ class TokenCommNonDeterministic(KernelComponent): - raise ValueError("sf_padding_block must be a multiple of 128.") +@@ -1391,7 +1391,7 @@ class TokenCommNonDeterministic(KernelComponent): + raise ValueError("token_in_stage_split must cut a token row into whole 16-byte chunks.") if self.token_back_enabled and self.combine_stage_depth > 1: - if self.token_back_schedule_mode != "static": @@ -66,7 +66,7 @@ index 7e2d5b2a70509a1a960fa450f95642ef321c9913..910e753adb0fcff4969ed89903d4d250 raise ValueError("Multi-stage token-back requires static scheduling.") if self.combine_format.is_quantized: raise ValueError("Multi-stage token-back supports BF16 combine only.") -@@ -1322,7 +1322,7 @@ class TokenCommNonDeterministic(KernelComponent): +@@ -1438,7 +1438,7 @@ class TokenCommNonDeterministic(KernelComponent): @property def token_back_push_data(self) -> bool: @@ -75,7 +75,7 @@ index 7e2d5b2a70509a1a960fa450f95642ef321c9913..910e753adb0fcff4969ed89903d4d250 @property def token_back_push_sf(self) -> bool: -@@ -1540,7 +1540,7 @@ class TokenCommNonDeterministic(KernelComponent): +@@ -1663,7 +1663,7 @@ class TokenCommNonDeterministic(KernelComponent): if self.token_back_enabled: fc2_done_slots = ( self.max_fc1_ready_slot_count @@ -84,7 +84,7 @@ index 7e2d5b2a70509a1a960fa450f95642ef321c9913..910e753adb0fcff4969ed89903d4d250 else self.experts_per_rank ) workspace.register( -@@ -1564,7 +1564,7 @@ class TokenCommNonDeterministic(KernelComponent): +@@ -1687,7 +1687,7 @@ class TokenCommNonDeterministic(KernelComponent): buffer_space="local", byte_alignment=128, ) @@ -93,7 +93,7 @@ index 7e2d5b2a70509a1a960fa450f95642ef321c9913..910e753adb0fcff4969ed89903d4d250 workspace.register( self.token_back_schedule_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" ) -@@ -1609,7 +1609,7 @@ class TokenCommNonDeterministic(KernelComponent): +@@ -1732,7 +1732,7 @@ class TokenCommNonDeterministic(KernelComponent): self._register_token_in_staging(transfer_overlay.add_lifetime("token_in")) if not self.token_back_enabled: return @@ -102,7 +102,7 @@ index 7e2d5b2a70509a1a960fa450f95642ef321c9913..910e753adb0fcff4969ed89903d4d250 back_overlay = workspace.create_overlay("nvlink.token_comm.main_smem.standalone_token_back") else: back_overlay = transfer_overlay -@@ -1664,20 +1664,20 @@ class TokenCommNonDeterministic(KernelComponent): +@@ -1788,7 +1788,7 @@ class TokenCommNonDeterministic(KernelComponent): combine_rings = self.combine_warp_count combine_depth = self.combine_stage_depth sf_padded = self.activation_sf_hidden_padded @@ -110,7 +110,8 @@ index 7e2d5b2a70509a1a960fa450f95642ef321c9913..910e753adb0fcff4969ed89903d4d250 + if self.token_back_mode == "standalone_warps": # nosec B105 available_bytes_per_warp = self.token_back_chunk_bytes else: - available_bytes_per_warp = self.bytes_per_token + sf_padded * int(self.activation_sf_dtype.width) // 8 + sf_bytes = sf_padded * int(self.activation_sf_dtype.width) // 8 if self.activation_is_quantized else 0 +@@ -1796,13 +1796,13 @@ class TokenCommNonDeterministic(KernelComponent): if self.token_back_push_data: bytes_per_output_token = self.hidden_size * int(self.combine_format.act_dtype.width) // 8 @@ -126,7 +127,7 @@ index 7e2d5b2a70509a1a960fa450f95642ef321c9913..910e753adb0fcff4969ed89903d4d250 raise ValueError("Token-back data chunk bytes must divide one row.") chunk_elements = chunk_bytes * 8 // int(self.combine_format.act_dtype.width) shape = ( -@@ -2176,7 +2176,7 @@ class TokenCommNonDeterministic(KernelComponent): +@@ -2303,7 +2303,7 @@ class TokenCommNonDeterministic(KernelComponent): ) token_in_size_barrier.arrive_and_wait() iket.range_pop() @@ -135,7 +136,7 @@ index 7e2d5b2a70509a1a960fa450f95642ef321c9913..910e753adb0fcff4969ed89903d4d250 sizes_ready_barrier = pipeline.NamedBarrier( barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count, -@@ -2335,7 +2335,7 @@ class TokenCommNonDeterministic(KernelComponent): +@@ -2495,7 +2495,7 @@ class TokenCommNonDeterministic(KernelComponent): flag_tracker.fire() cute.arch.sync_warp() @@ -144,7 +145,7 @@ index 7e2d5b2a70509a1a960fa450f95642ef321c9913..910e753adb0fcff4969ed89903d4d250 iket.range_push("token_in.transfer_barrier") transfer_lifetime_barrier = pipeline.NamedBarrier( barrier_id=self.transfer_lifetime_barrier_id, -@@ -2461,7 +2461,7 @@ class TokenCommNonDeterministic(KernelComponent): +@@ -2621,7 +2621,7 @@ class TokenCommNonDeterministic(KernelComponent): cute.arch.mbarrier_init(worker_mbarrier + Int32(init_stage), 1) cute.arch.sync_warp() @@ -153,7 +154,7 @@ index 7e2d5b2a70509a1a960fa450f95642ef321c9913..910e753adb0fcff4969ed89903d4d250 sizes_ready_barrier = pipeline.NamedBarrier( barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count, -@@ -2659,7 +2659,7 @@ class TokenCommNonDeterministic(KernelComponent): +@@ -2819,7 +2819,7 @@ class TokenCommNonDeterministic(KernelComponent): cute.arch.mbarrier_init(worker_mbarrier, 1) cute.arch.sync_warp() @@ -162,7 +163,7 @@ index 7e2d5b2a70509a1a960fa450f95642ef321c9913..910e753adb0fcff4969ed89903d4d250 sizes_ready_barrier = pipeline.NamedBarrier( barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count, -@@ -2675,7 +2675,7 @@ class TokenCommNonDeterministic(KernelComponent): +@@ -2835,7 +2835,7 @@ class TokenCommNonDeterministic(KernelComponent): activation_chunk_count = ceil_div(output_token_bytes, activation_chunk_bytes) data_window_unit = activation_chunk_bytes * 2 reuse_data_pacing_enabled = ( @@ -171,7 +172,7 @@ index 7e2d5b2a70509a1a960fa450f95642ef321c9913..910e753adb0fcff4969ed89903d4d250 ) # Preserve the empirical low:initial:high ratio of 1:2.5:5. data_average_window = Int32(data_window_unit) -@@ -2687,19 +2687,19 @@ class TokenCommNonDeterministic(KernelComponent): +@@ -2847,19 +2847,19 @@ class TokenCommNonDeterministic(KernelComponent): sf_chunk_count = ceil_div(output_sf_bytes, sf_chunk_bytes) sf_window_unit = ceil_div(sf_chunk_bytes * 2, 3) reuse_sf_pacing_enabled = ( @@ -194,7 +195,7 @@ index 7e2d5b2a70509a1a960fa450f95642ef321c9913..910e753adb0fcff4969ed89903d4d250 expert_ready_slot_begin = Int32(0) last_ready_slot = Int32(-1) -@@ -2707,7 +2707,7 @@ class TokenCommNonDeterministic(KernelComponent): +@@ -2867,7 +2867,7 @@ class TokenCommNonDeterministic(KernelComponent): while local_expert < Int32(self.experts_per_rank): expert_token_count = owned_sizes[local_expert] expert_valid_end = expert_valid_begin + expert_token_count @@ -203,7 +204,7 @@ index 7e2d5b2a70509a1a960fa450f95642ef321c9913..910e753adb0fcff4969ed89903d4d250 if next_dense_token < expert_valid_end: token_tile_count = (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) // Int32( self.tokens_per_fc1_ready_slot -@@ -2720,7 +2720,7 @@ class TokenCommNonDeterministic(KernelComponent): +@@ -2880,7 +2880,7 @@ class TokenCommNonDeterministic(KernelComponent): while next_dense_token < expert_valid_end: token_in_expert = next_dense_token - expert_valid_begin @@ -212,7 +213,7 @@ index 7e2d5b2a70509a1a960fa450f95642ef321c9913..910e753adb0fcff4969ed89903d4d250 # Same rank-local expert/tile prefix as FC1 ready; separate counter storage. ready_slot = expert_ready_slot_begin + token_in_expert // Int32(self.tokens_per_fc1_ready_slot) if ready_slot != last_ready_slot: -@@ -2765,7 +2765,7 @@ class TokenCommNonDeterministic(KernelComponent): +@@ -2925,7 +2925,7 @@ class TokenCommNonDeterministic(KernelComponent): current_chunk_bytes = Int32(chunk_bytes_this_round) current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) stateless_data_pacing_enabled = ( @@ -221,7 +222,7 @@ index 7e2d5b2a70509a1a960fa450f95642ef321c9913..910e753adb0fcff4969ed89903d4d250 and current_window_unit > self.minimum_pacing_window_cycles ) round_start_clock = Int64(0) -@@ -2841,7 +2841,7 @@ class TokenCommNonDeterministic(KernelComponent): +@@ -3001,7 +3001,7 @@ class TokenCommNonDeterministic(KernelComponent): current_chunk_bytes = Int32(chunk_bytes_this_round) current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) stateless_sf_pacing_enabled = ( @@ -230,7 +231,7 @@ index 7e2d5b2a70509a1a960fa450f95642ef321c9913..910e753adb0fcff4969ed89903d4d250 and current_window_unit > self.minimum_pacing_window_cycles ) round_start_clock = Int64(0) -@@ -2894,7 +2894,7 @@ class TokenCommNonDeterministic(KernelComponent): +@@ -3054,7 +3054,7 @@ class TokenCommNonDeterministic(KernelComponent): cute.arch.sync_warp() next_dense_token = self.next_token(next_dense_token) @@ -239,7 +240,7 @@ index 7e2d5b2a70509a1a960fa450f95642ef321c9913..910e753adb0fcff4969ed89903d4d250 # Advance even when this worker skipped the expert (empty experts contribute zero). expert_ready_slot_begin = expert_ready_slot_begin + ( expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1) -@@ -2910,9 +2910,9 @@ class TokenCommNonDeterministic(KernelComponent): +@@ -3070,9 +3070,9 @@ class TokenCommNonDeterministic(KernelComponent): # communication CTAs' combine warps, not every persistent CTA's transfer warps. global_worker_count = self.combine_cta_count * self.combine_warp_count schedule_counter = None @@ -265,10 +266,10 @@ index e9cf01a60538d05649f80d79c0cf8cac129566f7..ecc713fcbf9df7cba4ca999b45eee373 if not token_back_on_compute_ctas: dispatch_end = self.dispatch_warp_start + self.dispatch_warp_count diff --git a/communication/nvlink_domain/token_comm_deterministic.py b/communication/nvlink_domain/token_comm_deterministic.py -index 8b74d44c61c2c1efb3108309c49221e2c2b2588a..af3e8b6545b2c86aafa9ceb2a138cefe41d3d954 100644 +index 5f0f38f79326d591d9f1103bcc182cb95729c28f..6838edb87603bd90c56c76a47114582d5baf869a 100644 --- a/communication/nvlink_domain/token_comm_deterministic.py +++ b/communication/nvlink_domain/token_comm_deterministic.py -@@ -137,7 +137,7 @@ class _MetadataPushRouter(KernelComponent): +@@ -138,7 +138,7 @@ class _MetadataPushRouter(KernelComponent): sorted_metadata_region = "nvlink.token_comm.sorted_metadata" sorted_scores_region = "nvlink.token_comm.sorted_scores" pool_expert_base_region = "nvlink.token_comm.pool_expert_base" @@ -277,7 +278,7 @@ index 8b74d44c61c2c1efb3108309c49221e2c2b2588a..af3e8b6545b2c86aafa9ceb2a138cefe fc1_topk_scores_region = "nvlink.token_comm.fc1_topk_scores" source_expert_base_region = "nvlink.token_comm.source_expert_base" push_destination_base_region = "nvlink.token_comm.push_destination_base" -@@ -1347,15 +1347,15 @@ class TokenCommDeterministic(KernelComponent): +@@ -1375,15 +1375,15 @@ class TokenCommDeterministic(KernelComponent): fc2_activation_sf_region = "nvlink.token_comm.fc2_activation_sf" pre_reduced_activation_region = "nvlink.token_comm.pre_reduced_activation" pre_reduced_activation_sf_region = "nvlink.token_comm.pre_reduced_activation_sf" @@ -300,7 +301,7 @@ index 8b74d44c61c2c1efb3108309c49221e2c2b2588a..af3e8b6545b2c86aafa9ceb2a138cefe @classmethod def problem_desc_require(cls) -> dict[str, type]: -@@ -1519,7 +1519,7 @@ class TokenCommDeterministic(KernelComponent): +@@ -1573,7 +1573,7 @@ class TokenCommDeterministic(KernelComponent): @property def token_back_push_data(self) -> bool: @@ -309,7 +310,7 @@ index 8b74d44c61c2c1efb3108309c49221e2c2b2588a..af3e8b6545b2c86aafa9ceb2a138cefe @property def token_back_push_sf(self) -> bool: -@@ -1716,7 +1716,7 @@ class TokenCommDeterministic(KernelComponent): +@@ -1776,7 +1776,7 @@ class TokenCommDeterministic(KernelComponent): buffer_space="local", byte_alignment=128, ) @@ -318,7 +319,7 @@ index 8b74d44c61c2c1efb3108309c49221e2c2b2588a..af3e8b6545b2c86aafa9ceb2a138cefe workspace.register( self.token_back_schedule_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" ) -@@ -1764,14 +1764,14 @@ class TokenCommDeterministic(KernelComponent): +@@ -1825,14 +1825,14 @@ class TokenCommDeterministic(KernelComponent): if not self.token_back_enabled: return @@ -335,7 +336,7 @@ index 8b74d44c61c2c1efb3108309c49221e2c2b2588a..af3e8b6545b2c86aafa9ceb2a138cefe available_bytes_per_warp = self.standalone_chunk_bytes else: activation_bytes = self.bytes_per_token -@@ -1780,13 +1780,13 @@ class TokenCommDeterministic(KernelComponent): +@@ -1845,13 +1845,13 @@ class TokenCommDeterministic(KernelComponent): if self.token_back_push_data: bytes_per_output_token = self.hidden_size * int(self.combine_format.act_dtype.width) // 8 @@ -351,7 +352,7 @@ index 8b74d44c61c2c1efb3108309c49221e2c2b2588a..af3e8b6545b2c86aafa9ceb2a138cefe raise ValueError("Token-back data chunk bytes must divide one row.") chunk_elements = chunk_bytes * 8 // int(self.combine_format.act_dtype.width) token_back_lifetime.register_tensor( -@@ -2098,7 +2098,7 @@ class TokenCommDeterministic(KernelComponent): +@@ -2175,7 +2175,7 @@ class TokenCommDeterministic(KernelComponent): ) token_in_size_barrier.arrive_and_wait() iket.range_pop() @@ -360,7 +361,7 @@ index 8b74d44c61c2c1efb3108309c49221e2c2b2588a..af3e8b6545b2c86aafa9ceb2a138cefe sizes_ready_barrier = pipeline.NamedBarrier( barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count ) -@@ -2242,7 +2242,7 @@ class TokenCommDeterministic(KernelComponent): +@@ -2364,7 +2364,7 @@ class TokenCommDeterministic(KernelComponent): flag_tracker.fire() cute.arch.sync_warp() iket.range_pop() @@ -369,7 +370,7 @@ index 8b74d44c61c2c1efb3108309c49221e2c2b2588a..af3e8b6545b2c86aafa9ceb2a138cefe iket.range_push("token_in.transfer_barrier") transfer_lifetime_barrier = pipeline.NamedBarrier( barrier_id=self.transfer_lifetime_barrier_id, num_threads=self.transfer_thread_count -@@ -2310,7 +2310,7 @@ class TokenCommDeterministic(KernelComponent): +@@ -2432,7 +2432,7 @@ class TokenCommDeterministic(KernelComponent): cute.arch.sync_warp() owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) @@ -378,7 +379,7 @@ index 8b74d44c61c2c1efb3108309c49221e2c2b2588a..af3e8b6545b2c86aafa9ceb2a138cefe sizes_ready_barrier = pipeline.NamedBarrier( barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count ) -@@ -2325,7 +2325,7 @@ class TokenCommDeterministic(KernelComponent): +@@ -2447,7 +2447,7 @@ class TokenCommDeterministic(KernelComponent): activation_chunk_count = ceil_div(output_token_bytes, activation_chunk_bytes) data_window_unit = ceil_div(activation_chunk_bytes * 2, 3) reuse_data_pacing_enabled = ( @@ -387,7 +388,7 @@ index 8b74d44c61c2c1efb3108309c49221e2c2b2588a..af3e8b6545b2c86aafa9ceb2a138cefe ) # Preserve the empirical low:initial:high ratio of 1:2.5:5. data_average_window = Int32(data_window_unit) -@@ -2338,14 +2338,14 @@ class TokenCommDeterministic(KernelComponent): +@@ -2460,14 +2460,14 @@ class TokenCommDeterministic(KernelComponent): sf_chunk_count = ceil_div(output_sf_bytes, sf_chunk_bytes) sf_window_unit = ceil_div(sf_chunk_bytes * 2, 3) reuse_sf_pacing_enabled = ( @@ -404,7 +405,7 @@ index 8b74d44c61c2c1efb3108309c49221e2c2b2588a..af3e8b6545b2c86aafa9ceb2a138cefe next_dense_token = Int32(0) next_dense_token = self.next_token(next_dense_token) expert_valid_begin = Int32(0) -@@ -2401,7 +2401,7 @@ class TokenCommDeterministic(KernelComponent): +@@ -2523,7 +2523,7 @@ class TokenCommDeterministic(KernelComponent): current_chunk_bytes = Int32(chunk_bytes_this_round) current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) stateless_data_pacing_enabled = ( @@ -413,7 +414,7 @@ index 8b74d44c61c2c1efb3108309c49221e2c2b2588a..af3e8b6545b2c86aafa9ceb2a138cefe and current_window_unit > self.minimum_pacing_window_cycles ) round_start_clock = Int64(0) -@@ -2468,7 +2468,7 @@ class TokenCommDeterministic(KernelComponent): +@@ -2590,7 +2590,7 @@ class TokenCommDeterministic(KernelComponent): current_chunk_bytes = Int32(chunk_bytes_this_round) current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) stateless_sf_pacing_enabled = ( @@ -422,7 +423,7 @@ index 8b74d44c61c2c1efb3108309c49221e2c2b2588a..af3e8b6545b2c86aafa9ceb2a138cefe and current_window_unit > self.minimum_pacing_window_cycles ) round_start_clock = Int64(0) -@@ -2520,9 +2520,9 @@ class TokenCommDeterministic(KernelComponent): +@@ -2642,9 +2642,9 @@ class TokenCommDeterministic(KernelComponent): def next_token(self, current_token: Int32) -> Int32: global_worker_count = self.promised_launchable_sm_count * self.transfer_warp_count schedule_counter = None @@ -448,10 +449,10 @@ index 6e4d0fe7d363b8554e5b2789240e555ccc2499ea..ce4cb6465a0fb430bcb2a93acf2d8a15 # == 0) is read as packed b32 words and N fp32 values are written into # ``fp32_reg`` in code order. The 16 e2m1 values are exact in fp32, so the two diff --git a/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py b/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py -index cedc19ea477a24fd5d58e14777752fe84e44bd85..8057166cabff92d158798311a3fc7b1549d408b7 100644 +index a30dc98d4dd1bbae7600a6d3894fed8e7eec7c18..a363fcdd831a043315607dc277d0345706f80e6a 100644 --- a/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py +++ b/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py -@@ -959,6 +959,8 @@ class SwapABGatedActEpilogue(KernelComponent): +@@ -968,6 +968,8 @@ class SwapABGatedActEpilogue(KernelComponent): "intermediate_gateup_size": StaticOrRuntimeIntegerType, "combine_format": CombineFormat, "gate_up_clamp": Optional[float], @@ -460,7 +461,7 @@ index cedc19ea477a24fd5d58e14777752fe84e44bd85..8057166cabff92d158798311a3fc7b15 # SiTU (Kimi K3) selects a different gated-activation core; optional so # existing SwiGLU descriptors stay valid unchanged. See _resolve_situ_betas. "situ_beta": OptionalRequirement(Optional[float]), -@@ -999,6 +1001,8 @@ class SwapABGatedActEpilogue(KernelComponent): +@@ -1018,6 +1020,8 @@ class SwapABGatedActEpilogue(KernelComponent): raise ValueError(f"SiTU beta parameters must be positive, got {beta} and {linear_beta}.") if problem_desc["gate_up_clamp"] is not None: raise ValueError("SiTU does not support gate_up_clamp; the two activation variants are exclusive.") @@ -468,8 +469,17 @@ index cedc19ea477a24fd5d58e14777752fe84e44bd85..8057166cabff92d158798311a3fc7b15 + raise ValueError("SiTU does not support SwiGLUBias constants.") return float(beta), float(linear_beta) + @classmethod +@@ -1037,6 +1041,8 @@ class SwapABGatedActEpilogue(KernelComponent): + ) + if gate_activation != "swiglu" and problem_desc.get("situ_beta") is not None: + raise ValueError(f"gate_activation={gate_activation!r} cannot be combined with SiTU betas.") ++ if gate_activation != "swiglu" and problem_desc.get("swiglu_alpha") is not None: ++ raise ValueError(f"gate_activation={gate_activation!r} cannot be combined with SwiGLUBias constants.") + return gate_activation + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: -@@ -1010,6 +1014,13 @@ class SwapABGatedActEpilogue(KernelComponent): +@@ -1048,6 +1054,13 @@ class SwapABGatedActEpilogue(KernelComponent): self.intermediate_gateup_size = problem_desc["intermediate_gateup_size"] self.combine_format = problem_desc["combine_format"] self.gate_up_clamp = problem_desc["gate_up_clamp"] @@ -481,9 +491,9 @@ index cedc19ea477a24fd5d58e14777752fe84e44bd85..8057166cabff92d158798311a3fc7b15 + self.swiglu_alpha = float(self.swiglu_alpha) + self.swiglu_beta = float(self.swiglu_beta) self.situ_beta, self.situ_linear_beta = self._resolve_situ_betas(problem_desc) - self.mma_tiler_mnk = impl_desc["mma_tiler_mnk"] - self.cluster_shape_mn = impl_desc["cluster_shape_mn"] -@@ -1765,7 +1776,8 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): + self.gate_activation = self._resolve_gate_activation(problem_desc) + # The scratch costs mainloop SMEM budget, so only a descriptor that enables the scale +@@ -2021,7 +2034,8 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): # gate = min(gate, +limit) (upper bound only) # up = clamp(up, -limit, +limit) (symmetric) # 3. gated activation, either SwiGLU or SiTU (mutually exclusive; SiTU has no clamp): @@ -493,21 +503,20 @@ index cedc19ea477a24fd5d58e14777752fe84e44bd85..8057166cabff92d158798311a3fc7b15 # SiTU: out = beta*tanh(gate/beta)*sigmoid(gate) * linear_beta*tanh(up/linear_beta) # sigmoid(x) = rcp(1 + exp2(-x * log2e)) # -@@ -1777,8 +1789,12 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): - out = cute.make_rmem_tensor((n,), cutlass.Float32) - log2_e = 1.4426950408889634 - -- neg_log2e_pair = (cutlass.Float32(-log2_e), cutlass.Float32(-log2_e)) -+ activation_alpha = 1.0 if self.swiglu_alpha is None else self.swiglu_alpha -+ neg_activation_alpha_log2e = -activation_alpha * log2_e -+ neg_log2e_pair = (cutlass.Float32(neg_activation_alpha_log2e), cutlass.Float32(neg_activation_alpha_log2e)) +@@ -2037,8 +2051,12 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): + # (one Python-side double product, rounded once to fp32), so the SwiGLU instruction sequence + # is reused unchanged. SiTU never reaches here with quick_geglu (_resolve_gate_activation). + sigmoid_scale = self._quick_geglu_sigmoid_scale if self.gate_activation == "quick_geglu" else 1.0 ++ if cutlass.const_expr(self.swiglu_alpha is not None): ++ sigmoid_scale = self.swiglu_alpha + neg_log2e_pair = (cutlass.Float32(-sigmoid_scale * log2_e), cutlass.Float32(-sigmoid_scale * log2_e)) one_pair = (cutlass.Float32(1.0), cutlass.Float32(1.0)) + if cutlass.const_expr(self.swiglu_beta is not None): + beta_pair = (cutlass.Float32(self.swiglu_beta), cutlass.Float32(self.swiglu_beta)) if cutlass.const_expr(self.gate_up_clamp is not None): limit = cutlass.Float32(self.gate_up_clamp) - -@@ -1854,7 +1870,10 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): + if cutlass.const_expr(dequant_scale is not None and dequant_scales is not None): +@@ -2131,7 +2149,10 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): sigmoid_pair = _sigmoid(g0, g1) if cutlass.const_expr(self.situ_beta is None): @@ -520,7 +529,7 @@ index cedc19ea477a24fd5d58e14777752fe84e44bd85..8057166cabff92d158798311a3fc7b15 out_pair = cute.arch.mul_packed_f32x2(ug, sigmoid_pair) else: diff --git a/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py b/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py -index d7f4e53daff8e7e2cee8444b21c57f114c9a8fe1..7dad5c86e02b22f49021cb39c600d41f65243cd1 100644 +index ae93a57df562296bc79d08c4594c8c3cd9ba09aa..2e0497d3aa720edd03717fd0117fb5fd8d735499 100644 --- a/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py +++ b/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py @@ -73,6 +73,8 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): @@ -531,8 +540,8 @@ index d7f4e53daff8e7e2cee8444b21c57f114c9a8fe1..7dad5c86e02b22f49021cb39c600d41f + "swiglu_beta": OptionalRequirement(Optional[float]), "situ_beta": OptionalRequirement(Optional[float]), "situ_linear_beta": OptionalRequirement(Optional[float]), - "world_size": int, -@@ -127,6 +129,8 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): + # "swiglu" (default) or "quick_geglu"; validated by the epilogue, keyed in name(). +@@ -134,6 +136,8 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): self.b_major_mode = problem_desc["b_major_mode"] self.combine_format = problem_desc["combine_format"] self.gate_up_clamp = problem_desc["gate_up_clamp"] @@ -540,8 +549,8 @@ index d7f4e53daff8e7e2cee8444b21c57f114c9a8fe1..7dad5c86e02b22f49021cb39c600d41f + self.swiglu_beta = problem_desc.get("swiglu_beta") self.situ_beta = problem_desc.get("situ_beta") self.situ_linear_beta = problem_desc.get("situ_linear_beta") - self.world_size = problem_desc["world_size"] -@@ -168,8 +172,8 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): + self.gate_activation = problem_desc.get("gate_activation") or "swiglu" +@@ -177,8 +181,8 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): # static M is required only when an S>0 helper boundary must be # compiled into the scheduler gate. self.home_expert_count = self.local_expert_count @@ -552,7 +561,7 @@ index d7f4e53daff8e7e2cee8444b21c57f114c9a8fe1..7dad5c86e02b22f49021cb39c600d41f self.mixed_cga_config = NonClcMixedCgaConfig( preferred_cluster_shape=self.cluster_shape_mn, -@@ -453,6 +457,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): +@@ -463,6 +467,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): f"epiflag{epi_flags}_clusters{self.launch_cluster_count}_ctas{self.mixed_cga_config.total_cta_cnt}_" f"tokeninflag{self.token_in_flag_batch}_tokenback{self.token_back_mode}_" f"combine{self.combine_format.name}_clamp{self.gate_up_clamp}_" @@ -560,7 +569,7 @@ index d7f4e53daff8e7e2cee8444b21c57f114c9a8fe1..7dad5c86e02b22f49021cb39c600d41f # The SiTU betas are codegen-time constants, so they must key the compiled kernel. f"situ{self.situ_beta}x{self.situ_linear_beta}_" f"{'apply_topk_fc1' if self.apply_topk_at_fc1 else 'apply_topk_fc2'}_" -@@ -1032,7 +1037,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): +@@ -1084,7 +1089,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): pipeline_init_wait_mixed_cga( self.cluster_shape_mn, self.resolved_fallback_cluster_shape_mn, is_fallback_cluster ) @@ -569,7 +578,7 @@ index d7f4e53daff8e7e2cee8444b21c57f114c9a8fe1..7dad5c86e02b22f49021cb39c600d41f finalize_barrier = pipeline.NamedBarrier(barrier_id=13, num_threads=finalize_barrier_thread_count) if warp_idx == self.scheduler_warp_id: -@@ -1161,7 +1166,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): +@@ -1213,7 +1218,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): self.token_comm.token_in(self._smem_workspace, smem_base) iket.range_pop() if cutlass.const_expr( @@ -578,7 +587,7 @@ index d7f4e53daff8e7e2cee8444b21c57f114c9a8fe1..7dad5c86e02b22f49021cb39c600d41f ): iket.range_push("mega.token_back") self.token_comm.token_back(self._smem_workspace, smem_base) -@@ -1170,7 +1175,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): +@@ -1222,7 +1227,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): iket.range_push("mega.tail_reset") self.token_comm.reset_tail() iket.range_pop() @@ -588,7 +597,7 @@ index d7f4e53daff8e7e2cee8444b21c57f114c9a8fe1..7dad5c86e02b22f49021cb39c600d41f self.token_comm.token_back(self._smem_workspace, smem_base) iket.range_pop() diff --git a/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py b/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py -index 03507e3bb60f159fca2343204339810a3db2d14d..7fbb2be22b7645c0fe116b4db662c64a053bda65 100644 +index a97e10dd651fd728fd1e73dd3ec2791ee1b4ef9a..911ced925fb0d789d19983ec68e817cb36a0c329 100644 --- a/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py +++ b/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py @@ -63,6 +63,13 @@ class SwapABGatedActEpilogue(_SharedSwapABGatedActEpilogue): @@ -603,8 +612,8 @@ index 03507e3bb60f159fca2343204339810a3db2d14d..7fbb2be22b7645c0fe116b4db662c64a + self.swiglu_alpha = float(self.swiglu_alpha) + self.swiglu_beta = float(self.swiglu_beta) self.situ_beta, self.situ_linear_beta = self._resolve_situ_betas(problem_desc) - self.mma_tiler_mnk = impl_desc["mma_tiler_mnk"] - self.cluster_shape_mn = impl_desc["cluster_shape_mn"] + # quick_geglu is Blackwell-only so far; refuse rather than silently computing SwiGLU. + if problem_desc.get("gate_activation") not in (None, "swiglu"): diff --git a/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_local_mega_moe_kernel.py b/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_local_mega_moe_kernel.py index a486110e0e3eb4b673d5857f7d08484cc24893a5..3e86b1ad34e2e31ab8b3a394208e921156c4b552 100644 --- a/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_local_mega_moe_kernel.py @@ -661,7 +670,7 @@ index ce7877ec231429f6487668527285388c09bea0bd..fd3ceb087e89c02c041b0e794dc22fbc @classmethod def problem_desc_require(cls) -> dict[str, type]: diff --git a/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py b/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py -index bd269d76d534d195778186d3844982101ee87860..f0868f743ba480f8d3c8f0a9aac7bf31cdc90daa 100644 +index 21f00a540720eff436d738b706819ef85502f023..2d503d43c52ff7ef9e04a00a6dc05c993b07662f 100644 --- a/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py +++ b/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py @@ -896,6 +896,8 @@ class SwapABGatedActEpilogue(KernelComponent): @@ -694,9 +703,9 @@ index bd269d76d534d195778186d3844982101ee87860..f0868f743ba480f8d3c8f0a9aac7bf31 + self.swiglu_alpha = float(self.swiglu_alpha) + self.swiglu_beta = float(self.swiglu_beta) self.situ_beta, self.situ_linear_beta = self._resolve_situ_betas(problem_desc) - self.mma_tiler_mnk = impl_desc["mma_tiler_mnk"] - self.cluster_shape_mn = impl_desc["cluster_shape_mn"] -@@ -1639,7 +1650,8 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): + # The Blackwell epilogue accepts gate_activation="quick_geglu"; this one has not been + # ported yet. Refuse rather than silently computing SwiGLU for a quick_geglu descriptor. +@@ -1645,7 +1656,8 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): # gate = min(gate, +limit) (upper bound only) # up = clamp(up, -limit, +limit) (symmetric) # 3. gated activation, either SwiGLU or SiTU (mutually exclusive; SiTU has no clamp): @@ -706,7 +715,7 @@ index bd269d76d534d195778186d3844982101ee87860..f0868f743ba480f8d3c8f0a9aac7bf31 # SiTU: out = beta*tanh(gate/beta)*sigmoid(gate) * linear_beta*tanh(up/linear_beta) # sigmoid(x) = rcp(1 + exp2(-x * log2e)) # -@@ -1651,8 +1663,12 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): +@@ -1657,8 +1669,12 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): out = cute.make_rmem_tensor((n,), cutlass.Float32) log2_e = 1.4426950408889634 @@ -720,7 +729,7 @@ index bd269d76d534d195778186d3844982101ee87860..f0868f743ba480f8d3c8f0a9aac7bf31 if cutlass.const_expr(self.gate_up_clamp is not None): limit = cutlass.Float32(self.gate_up_clamp) -@@ -1728,7 +1744,10 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): +@@ -1734,7 +1750,10 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): sigmoid_pair = _sigmoid(g0, g1) if cutlass.const_expr(self.situ_beta is None): @@ -732,7 +741,7 @@ index bd269d76d534d195778186d3844982101ee87860..f0868f743ba480f8d3c8f0a9aac7bf31 ug = cute.arch.mul_packed_f32x2((u0, u1), (g0, g1)) out_pair = cute.arch.mul_packed_f32x2(ug, sigmoid_pair) else: -@@ -1985,7 +2004,7 @@ class SwapABFc2Epilogue(_ImmutableAfterInit): +@@ -1991,7 +2010,7 @@ class SwapABFc2Epilogue(_ImmutableAfterInit): publish: cutlass.Constexpr = self.token_back_enabled flag_address = Int64(0) if cutlass.const_expr(publish): @@ -853,10 +862,10 @@ index f49de382daeef5a9751a19dd190008f02827c54c..3c5bb77f999cd0013dd68ec11668adef self.token_comm.token_back(self._smem_workspace, smem_base) iket.range_pop() diff --git a/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py b/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py -index 96b8342673f6943862b81bd9ea8e3aa85d1fac30..c5567f3df5a0031dd77a9c038d9b5b03f510ebd8 100644 +index 043f2f86d7e9f0017ca28144f3d56b85eac4aa7b..9d99afc3344a55a00544fddbe3979dd0b0bf3032 100644 --- a/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py +++ b/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py -@@ -109,6 +109,8 @@ class BlockScaledSwapAbGenphaseMoeKernel(KernelClass): +@@ -108,6 +108,8 @@ class BlockScaledSwapAbGenphaseMoeKernel(KernelClass): "quant_kind": str, "combine_format": CombineFormat, "gate_up_clamp": Optional[float], @@ -865,7 +874,7 @@ index 96b8342673f6943862b81bd9ea8e3aa85d1fac30..c5567f3df5a0031dd77a9c038d9b5b03 "world_size": int, "topk": int, "topk_index_dtype": type, -@@ -169,6 +171,8 @@ class BlockScaledSwapAbGenphaseMoeKernel(KernelClass): +@@ -168,6 +170,8 @@ class BlockScaledSwapAbGenphaseMoeKernel(KernelClass): self.acc_dtype = tcgen05_block_scaled_acc_dtype self.combine_format = problem_desc["combine_format"] self.gate_up_clamp = problem_desc["gate_up_clamp"] @@ -874,7 +883,7 @@ index 96b8342673f6943862b81bd9ea8e3aa85d1fac30..c5567f3df5a0031dd77a9c038d9b5b03 self.world_size = problem_desc["world_size"] self.topk = problem_desc["topk"] self.topk_index_dtype = problem_desc["topk_index_dtype"] -@@ -546,6 +550,7 @@ class BlockScaledSwapAbGenphaseMoeKernel(KernelClass): +@@ -545,6 +549,7 @@ class BlockScaledSwapAbGenphaseMoeKernel(KernelClass): f"epiflag{epi_flags}_clusters{self.launch_cluster_count}_" f"pusher{self.pusher_cta_count}_refine{self.refine_participant_groups}x{self.refine_output_stages}_" f"combine{self.combine_format.name}_clamp{self.gate_up_clamp}_" diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 23dfb7a6c9ac..23d66e54e12b 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -3,7 +3,7 @@ vendors: cutedsl-megamoe: url: https://gitlab-master.nvidia.com/bangyus/cutedsl_megamoe.git branch: main - commit: 382ce4a6f111d0769f31c4c975e1c643f0dc4f70 + commit: 7832a3300e24d791f94b09b3f6b37b5429e2fdf5 source: sources destination: tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe include: @@ -19,8 +19,8 @@ vendors: - kernel_src/schedulers/[!w]*.py - kernel_src/schedulers/work_id_claim.py patch: 3rdparty/vendor_patches/cutedsl-megamoe.patch - patch_digest: sha256:0676cec5fe7c087f93bb6aec9d15324460660e74f3caa94a86d7fcb7253deda8 - digest: sha256-tree-v1:bda2e41e83acec2801ecb5020b4fadbab50efb051c1911b5a204e5c0efdba96d + patch_digest: sha256:f24f6c5d296410bf3393ee5aafbb967982eee373a6e5f3289aed1b801b9b9d96 + digest: sha256-tree-v1:6e83535fe953947cc2b3c0e26a22cd8ded0dba1ff7845965e1f6b7ae36b2aac3 flashinfer-kda-bt16: url: https://github.com/flashinfer-ai/flashinfer.git branch: main diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm.py index 910e753adb0f..374aa01b055b 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm.py @@ -44,13 +44,23 @@ @dataclasses.dataclass(frozen=True) class TokenCommArgs: - """Actual tensor and peer-mapping values for one communication invocation.""" + """Device views materialized inside the fused kernel region. + + ``pre_reduced_activation`` is the per-topk combine staging plane: TokenComm's own symmetric region, except + under in-kernel top-k reduction where it degenerates to a view of the caller's 2D output. That REDG + accumulates, so in that mode the incoming content is the caller's accumulation base -- zero, or a + shared-expert result. + + ``fc1_activation_pool`` is the caller-owned ``(pool_capacity, hidden)`` dispatch pool that replaces + TokenComm's own ``fc1_activation_region`` under ``external_dispatch_pool``; ``None`` otherwise. + """ activation: cute.Tensor - activation_sf: cute.Tensor + activation_sf: Optional[cute.Tensor] pre_reduced_activation: cute.Tensor pre_reduced_activation_sf: Optional[cute.Tensor] peer_rank_ptr_mapper: SymmetricBufferDevice + fc1_activation_pool: Optional[cute.Tensor] = None @cute.jit @@ -115,6 +125,8 @@ class _MetadataPushRouter(KernelComponent): pool_expert_base_region = "nvlink.token_comm.pool_expert_base" token_src_metadata_region = "nvlink.token_comm.token_src_metadata" # nosec B105 fc1_topk_scores_region = "nvlink.token_comm.fc1_topk_scores" + # Pool-ordered fp32 per-token FC1 activation scale, pushed next to the metadata like the scores. + fc1_activation_scales_region = "nvlink.token_comm.fc1_activation_per_token_scales" source_expert_base_region = "nvlink.token_comm.source_expert_base" push_destination_base_region = "nvlink.token_comm.push_destination_base" sorted_metadata_ready_region = "nvlink.token_comm.sorted_metadata_ready" @@ -142,6 +154,8 @@ def problem_desc_require(cls) -> dict[str, type]: "topk": int, "max_tokens_per_rank": int, "apply_topk_at_fc1": bool, + # The FC1 activation was quantized from x / s_t; the router carries s_t to the pool. + "fc1_activation_per_token_scale": OptionalRequirement(bool), } @classmethod @@ -160,6 +174,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.topk = problem_desc["topk"] self.max_tokens_per_rank = problem_desc["max_tokens_per_rank"] self.apply_topk_at_fc1 = problem_desc["apply_topk_at_fc1"] + self.fc1_activation_per_token_scale = bool(problem_desc.get("fc1_activation_per_token_scale", False)) self.token_padding_block = impl_desc["token_padding_block"] self.promised_launchable_sm_count = impl_desc["promised_launchable_sm_count"] @@ -319,6 +334,13 @@ def _register_router_workspace(self, workspace: DeviceWorkspace) -> None: workspace.register( self.fc1_topk_scores_region, cutlass.Float32, (self.worst_case_token_count,), buffer_space="shared" ) + if self.fc1_activation_per_token_scale: + workspace.register( + self.fc1_activation_scales_region, + cutlass.Float32, + (self.worst_case_token_count,), + buffer_space="shared", + ) workspace.register(self.pool_expert_base_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local") workspace.register( self.source_expert_base_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="local" @@ -364,14 +386,18 @@ def launch_router( peer_rank_ptr_mapper_host, device_workspace: DeviceWorkspace, stream: cuda.CUstream, + activation_scales: Optional[cute.Tensor] = None, # (max_tokens_per_rank,) fp32, local ) -> None: """Launch counting-sort DATA, size-exchange HELPER, and metadata PUSH roles.""" if cutlass.const_expr(self.apply_topk_at_fc1 and topk_scores is None): raise ValueError("apply_topk_at_fc1 requires router topk_scores.") + if cutlass.const_expr(self.fc1_activation_per_token_scale != (activation_scales is not None)): + raise ValueError("The router carries activation_scales iff fc1_activation_per_token_scale is enabled.") peer_rank_ptr_mapper = peer_rank_ptr_mapper_host.make_device_object() self._router_kernel( topk_indices, topk_scores, + activation_scales, local_rank, local_workspace, shared_workspace, @@ -389,6 +415,7 @@ def _router_kernel( self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor], + activation_scales: Optional[cute.Tensor], local_rank: Int32, local_workspace: cute.Pointer, shared_workspace: cute.Pointer, @@ -417,13 +444,14 @@ def _router_kernel( self._router_grid_thread_idx = grid_thread_idx self._router_warp_idx = warp_idx self._router_lane_idx = lane_idx + self._router_activation_scales = activation_scales if cutlass.const_expr(self.router_data_cta_count == 1): self._router_single_cta(topk_indices, topk_scores, smem_base) else: self._router_multiple_ctas(topk_indices, topk_scores, smem_base) if linear_cta_idx < Int32(self.router_push_cta_count): - self._router_push_metadata() + self._router_push_metadata(smem_base) device_workspace.remove_device_members() self._device_workspace = None @@ -434,6 +462,7 @@ def _router_kernel( self._router_grid_thread_idx = None self._router_warp_idx = None self._router_lane_idx = None + self._router_activation_scales = None @cute.jit def _router_single_cta( @@ -646,7 +675,26 @@ def _router_helper_multiple_ctas(self, smem_base: cute.Pointer) -> None: iket.range_pop() @cute.jit - def _router_push_metadata(self) -> None: + def _expert_of_route(self, source_expert_base: cute.Tensor, route: Int32) -> Int32: + """Expert whose segment of this rank's sorted route array contains ``route``. + + ``source_expert_base`` is the exclusive prefix of this rank's per-expert route counts (non-decreasing), so the + answer is the largest expert with ``source_expert_base[expert] <= route`` (upper_bound - 1). A zero-size expert + shares its base with its successor and is therefore never returned. Binary search over the padded expert + count: 256 experts -> 8 steps, each one smem load. + """ + search_steps = (self.expert_count_padded - 1).bit_length() + expert = Int32(0) + for step_idx in cutlass.range_constexpr(search_steps): + step = Int32(1 << (search_steps - 1 - step_idx)) + probe = expert + step + if probe < Int32(self.expert_count_padded): + if source_expert_base[probe] <= route: + expert = probe + return expert + + @cute.jit + def _router_push_metadata(self, smem_base: cute.Pointer) -> None: block_thread_count = self.router_warps_per_cta * 32 sorted_metadata_ready = self._device_workspace.ptr(self.sorted_metadata_ready_region) push_table_ready = self._device_workspace.ptr(self.push_table_ready_region) @@ -659,39 +707,75 @@ def _router_push_metadata(self) -> None: nanosleep(150) cute.arch.sync_threads() - global_expert = self._router_linear_cta_idx * Int32(self.router_warps_per_cta) + self._router_warp_idx - if global_expert < Int32(self.expert_count): - sizes_by_rank = self._device_workspace.tensor(self.sizes_by_rank_region) - source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) - push_destination_base = self._device_workspace.tensor(self.push_destination_base_region) - route_count = sizes_by_rank[self._router_local_rank, global_expert] - source_begin = source_expert_base[global_expert] - destination_begin = push_destination_base[global_expert] - destination_rank = global_expert // Int32(self.experts_per_rank) - peer_offset = self._peer_rank_ptr_mapper.map(Int64(0), destination_rank, Int64(0)) - source_metadata = self._device_workspace.ptr(self.sorted_metadata_region) - destination_metadata_address = ( - self._device_workspace.ptr(self.token_src_metadata_region).toint() + peer_offset - ) - if cutlass.const_expr(self.apply_topk_at_fc1): - source_scores = self._device_workspace.ptr(self.sorted_scores_region) - destination_scores_address = ( - self._device_workspace.ptr(self.fc1_topk_scores_region).toint() + peer_offset + # Partition the push by route index, not by expert: one warp per expert serialises that expert's whole route + # run (32 routes per latency-bound remote-store round), so under skewed routing the hottest expert's warp runs + # for as long as the whole kernel while the other warps idle. Every push warp owns an equal, 32-aligned slice + # of this rank's flat sorted route array and finds the owning expert with a binary search over + # source_expert_base (exclusive prefix, non-decreasing; the largest expert with base <= route wins, which skips + # zero-size experts). Completion signalling is unchanged: per CTA (below), not per expert. + sizes_by_rank = self._device_workspace.tensor(self.sizes_by_rank_region) + source_expert_base_gmem = self._device_workspace.tensor(self.source_expert_base_region) + push_destination_base_gmem = self._device_workspace.tensor(self.push_destination_base_region) + # Stage both tables in smem for the per-lane lookups below (the DATA-phase prefix/histogram regions are dead + # once sorted_metadata_ready == data_cta_count, which this CTA already waited for); the binary search would + # otherwise issue 8 dependent global loads per route. + source_expert_base = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) + push_destination_base = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) + table_rounds = ceil_div(self.expert_count_padded, block_thread_count) + for table_round in cutlass.range_constexpr(table_rounds): + table_expert = Int32(table_round * block_thread_count) + self._router_thread_idx + if table_expert < Int32(self.expert_count_padded): + source_expert_base[table_expert] = source_expert_base_gmem[table_expert] + push_destination_base[table_expert] = push_destination_base_gmem[table_expert] + cute.arch.sync_threads() + source_metadata = self._device_workspace.ptr(self.sorted_metadata_region) + destination_metadata_base = self._device_workspace.ptr(self.token_src_metadata_region).toint() + if cutlass.const_expr(self.apply_topk_at_fc1): + source_scores = self._device_workspace.ptr(self.sorted_scores_region) + destination_scores_base = self._device_workspace.ptr(self.fc1_topk_scores_region).toint() + if cutlass.const_expr(self.fc1_activation_per_token_scale): + destination_activation_scales_base = self._device_workspace.ptr(self.fc1_activation_scales_region).toint() + # This rank's routes are one flat array of total_routes records, sorted by expert; every push warp takes an + # equal, 32-aligned slice of it (the last warps may get an empty slice) so the work per warp is total_routes / + # push_warp_count regardless of how skewed the per-expert counts are. + last_expert = Int32(self.expert_count - 1) + total_routes = source_expert_base[last_expert] + sizes_by_rank[self._router_local_rank, last_expert] + push_warp_count = self.router_push_cta_count * self.router_warps_per_cta + push_warp_idx = self._router_linear_cta_idx * Int32(self.router_warps_per_cta) + self._router_warp_idx + # ceil(total_routes / push_warp_count), rounded up to a multiple of 32 (one record per lane per round). + routes_per_warp = ((total_routes + Int32(push_warp_count - 1)) // Int32(push_warp_count) + Int32(31)) // Int32(32) * Int32(32) + slice_begin = push_warp_idx * routes_per_warp + slice_end = cutlass.min(slice_begin + routes_per_warp, total_routes) + route_round_count = Int32(0) + if slice_begin < total_routes: + route_round_count = (slice_end - slice_begin + Int32(31)) // Int32(32) + for route_round in cutlass.range(route_round_count, unroll=1): + route = slice_begin + Int32(route_round) * Int32(32) + self._router_lane_idx + if route < slice_end: + # The record does not carry its expert; the slice may span several experts, so each lane looks up + # the expert whose segment contains its route and maps the in-segment offset to the destination. + expert = self._expert_of_route(source_expert_base, route) + destination_position = push_destination_base[expert] + (route - source_expert_base[expert]) + destination_rank = expert // Int32(self.experts_per_rank) + peer_offset = self._peer_rank_ptr_mapper.map(Int64(0), destination_rank, Int64(0)) + metadata = cute.arch.load(source_metadata + route, cutlass.Int64) + stg_b64( + destination_metadata_base + peer_offset + Int64(destination_position) * Int64(TokenSrcMetadata.nbytes), + metadata, ) - route_round_count = (route_count + Int32(31)) // Int32(32) - for route_round in cutlass.range(route_round_count, unroll=1): - route = Int32(route_round) * Int32(32) + self._router_lane_idx - if route < route_count: - source_position = source_begin + route - destination_position = destination_begin + route - metadata = cute.arch.load(source_metadata + source_position, cutlass.Int64) - stg_b64( - destination_metadata_address + Int64(destination_position) * Int64(TokenSrcMetadata.nbytes), - metadata, + if cutlass.const_expr(self.apply_topk_at_fc1): + score = cute.arch.load(source_scores + route, cutlass.Float32) + stg_f32(destination_scores_base + peer_offset + Int64(destination_position) * Int64(4), score) + if cutlass.const_expr(self.fc1_activation_per_token_scale): + # The scale is per source token, not per route, so it is not sorted: gather it + # here by the token the metadata already names (low 32 bits, see + # TokenSrcMetadata) and push it into the same pool slot as the score. + source_token = Int32(Int64(metadata) & Int64(0xFFFFFFFF)) + activation_scale = self._router_activation_scales[source_token] + stg_f32( + destination_activation_scales_base + peer_offset + Int64(destination_position) * Int64(4), + activation_scale, ) - if cutlass.const_expr(self.apply_topk_at_fc1): - score = cute.arch.load(source_scores + source_position, cutlass.Float32) - stg_f32(destination_scores_address + Int64(destination_position) * Int64(4), score) cute.arch.sync_threads() if self._router_thread_idx == Int32(0): @@ -1064,6 +1148,13 @@ def fc1_topk_scores_tensor(self, device_workspace: DeviceWorkspace) -> Optional[ return None return device_workspace.tensor(self.fc1_topk_scores_region) + @cute.jit + def fc1_activation_per_token_scales_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + """Pool-ordered (worst_case_token_count,) fp32 view, or None when the feature is off.""" + if cutlass.const_expr(not self.fc1_activation_per_token_scale): + return None + return device_workspace.tensor(self.fc1_activation_scales_region) + class TokenCommNonDeterministic(KernelComponent): """Public fused-kernel token communication component.""" @@ -1131,6 +1222,7 @@ def impl_desc_require(cls) -> dict[str, type]: "comm_stage_depths": OptionalRequirement(tuple), "token_back_chunk_bytes": OptionalRequirement(int), "dispatch_sf_group": OptionalRequirement(int), + "token_in_stage_split": OptionalRequirement(bool), } def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: @@ -1141,7 +1233,8 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.topk = problem_desc["topk"] self.max_tokens_per_rank = problem_desc["max_tokens_per_rank"] self.hidden_size = problem_desc["hidden_size"] - self.quant_kind = QuantKind(problem_desc["quant_kind"]) + self.quant_kind = "bf16" if problem_desc["quant_kind"] == "bf16" else QuantKind(problem_desc["quant_kind"]) + self.activation_quant_kind = None if self.quant_kind == "bf16" else self.quant_kind self.combine_format = problem_desc["combine_format"] self.apply_topk_at_fc1 = problem_desc["apply_topk_at_fc1"] @@ -1193,6 +1286,8 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.token_back_chunk_bytes = impl_desc.get("token_back_chunk_bytes", self.standalone_chunk_bytes) # Tokens a token_in warp batches before writing their scale rows; see _token_in_multi_stage. self.dispatch_sf_group = impl_desc.get("dispatch_sf_group", 1) + self.token_in_stage_split = impl_desc.get("token_in_stage_split", False) + self._token_in_stage_chunk_count = 2 if self.token_in_stage_split else 1 self._validate_configuration() self._router = _MetadataPushRouter(problem_desc, impl_desc) @@ -1274,11 +1369,26 @@ def _validate_configuration(self) -> None: raise ValueError("fc2_done_signals_per_token_tile must be positive when token-back is enabled.") if self.reduce_topk_in_kernel and self.combine_format.act_dtype is not cutlass.BFloat16: raise ValueError("In-kernel top-k reduction requires BF16 combine data.") - element_block = self.activation_sf_vector_size * 4 - if self.hidden_size % element_block != 0: - raise ValueError(f"{self.quant_kind} requires hidden_size divisible by {element_block}.") + if self.activation_is_quantized: + element_block = self.activation_sf_vector_size * 4 + if self.hidden_size % element_block != 0: + raise ValueError(f"{self.quant_kind} requires hidden_size divisible by {element_block}.") + else: + # The 1D bulk copies move whole token rows, so each row must be whole 16-byte lines. + row_element_block = 128 // int(self.activation_dtype.width) + if self.hidden_size % row_element_block != 0: + raise ValueError(f"The BF16 wire requires hidden_size divisible by {row_element_block}.") + if self.dispatch_stage_depth != 1: + raise NotImplementedError("The BF16 wire requires comm_stage_depths[0]=1.") if self.sf_padding_block % 128 != 0: raise ValueError("sf_padding_block must be a multiple of 128.") + if self.token_in_stage_split: + if self.activation_is_quantized or self.token_back_enabled: + raise NotImplementedError( + "token_in_stage_split requires the BF16 wire with token_back disabled (epi_warps)." + ) + if self.bytes_per_token % (self._token_in_stage_chunk_count * 16) != 0: + raise ValueError("token_in_stage_split must cut a token row into whole 16-byte chunks.") if self.token_back_enabled and self.combine_stage_depth > 1: if self.token_back_schedule_mode != "static": # nosec B105 @@ -1290,17 +1400,21 @@ def _validate_configuration(self) -> None: def experts_per_rank(self) -> int: return self.expert_count // self.world_size + @property + def activation_is_quantized(self) -> bool: + return self.activation_quant_kind is not None + @property def activation_dtype(self) -> type: - return self.quant_kind.activation_dtype + return cutlass.BFloat16 if self.activation_quant_kind is None else self.activation_quant_kind.activation_dtype @property - def activation_sf_dtype(self) -> type: - return self.quant_kind.sf_dtype + def activation_sf_dtype(self) -> Optional[type]: + return None if self.activation_quant_kind is None else self.activation_quant_kind.sf_dtype @property def activation_sf_vector_size(self) -> int: - return self.quant_kind.sf_vec_size + return 0 if self.activation_quant_kind is None else self.activation_quant_kind.sf_vec_size @property def bytes_per_token(self) -> int: @@ -1308,6 +1422,8 @@ def bytes_per_token(self) -> int: @property def activation_sf_hidden_padded(self) -> int: + if not self.activation_is_quantized: + return 0 valid_hidden = self.hidden_size // self.activation_sf_vector_size elements_per_16_bytes = 128 // int(self.activation_sf_dtype.width) return int(round_up(valid_hidden, elements_per_16_bytes)) @@ -1400,6 +1516,7 @@ def launch_router( peer_rank_ptr_mapper_host, device_workspace: DeviceWorkspace, stream: cuda.CUstream, + activation_scales: Optional[cute.Tensor] = None, ) -> None: self._router.launch_router( topk_indices, @@ -1410,6 +1527,7 @@ def launch_router( peer_rank_ptr_mapper_host, device_workspace, stream, + activation_scales, ) @cute.jit @@ -1428,6 +1546,10 @@ def token_src_metadata_tensor(self, device_workspace: DeviceWorkspace) -> cute.T def fc1_topk_scores_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: return self._router.fc1_topk_scores_tensor(device_workspace) + @cute.jit + def fc1_activation_per_token_scales_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + return self._router.fc1_activation_per_token_scales_tensor(device_workspace) + @cute.jit def assign_device_members( self, @@ -1528,14 +1650,15 @@ def _register_main_workspace(self, workspace: DeviceWorkspace) -> None: buffer_space="local", byte_alignment=128, ) - activation_sf_element_count = self.worst_case_sf_token_count * self.activation_sf_hidden_padded - workspace.register( - self.fc1_activation_sf_region, - self.activation_sf_dtype, - (activation_sf_element_count,), - buffer_space="local", - byte_alignment=128, - ) + if self.activation_is_quantized: + activation_sf_element_count = self.worst_case_sf_token_count * self.activation_sf_hidden_padded + workspace.register( + self.fc1_activation_sf_region, + self.activation_sf_dtype, + (activation_sf_element_count,), + buffer_space="local", + byte_alignment=128, + ) if self.token_back_enabled: fc2_done_slots = ( @@ -1621,7 +1744,7 @@ def _register_token_in_staging(self, staging_lifetime) -> None: dispatch_depth = self.dispatch_stage_depth sf_padded = self.activation_sf_hidden_padded if dispatch_depth == 1: - activation_shape = (dispatch_rings, self.hidden_size) + activation_shape = (dispatch_rings, self.hidden_size // self._token_in_stage_chunk_count) sf_shape = (dispatch_rings, (self.activation_sf_vector_size, sf_padded)) sf_stride = (sf_padded, (0, 1)) else: @@ -1638,14 +1761,15 @@ def _register_token_in_staging(self, staging_lifetime) -> None: activation_shape, byte_alignment=16, ) - staging_lifetime.register_tensor( - self.token_in_sf_smem_region, - self.activation_sf_dtype, - sf_shape, - stride=sf_stride, - byte_alignment=16, - ) - if self.dispatch_stage_depth > 1: + if self.activation_is_quantized: + staging_lifetime.register_tensor( + self.token_in_sf_smem_region, + self.activation_sf_dtype, + sf_shape, + stride=sf_stride, + byte_alignment=16, + ) + if self.activation_is_quantized and self.dispatch_stage_depth > 1: # One k-group of a token group, laid out the way the destination wants it. Shares the token_in # lifetime because it is only live between a group's first arrival and its scale store. The # condition is the depth rather than the group size, because it is the depth that selects the @@ -1667,7 +1791,8 @@ def _register_token_back_staging(self, staging_lifetime) -> None: if self.token_back_mode == "standalone_warps": # nosec B105 available_bytes_per_warp = self.token_back_chunk_bytes else: - available_bytes_per_warp = self.bytes_per_token + sf_padded * int(self.activation_sf_dtype.width) // 8 + sf_bytes = sf_padded * int(self.activation_sf_dtype.width) // 8 if self.activation_is_quantized else 0 + available_bytes_per_warp = self.bytes_per_token + sf_bytes if self.token_back_push_data: bytes_per_output_token = self.hidden_size * int(self.combine_format.act_dtype.width) // 8 @@ -1716,7 +1841,9 @@ def fc1_activation_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tenso ) @cute.jit - def fc1_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + def fc1_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.activation_is_quantized): + return None layout = tile_atom_to_shape_SF( (self.worst_case_sf_token_count, self.hidden_size, 1), self.activation_sf_vector_size ) @@ -2198,30 +2325,31 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No token_in_mbarriers = smem_workspace.ptr(self.token_in_mbarrier_region, smem_base) token_in_activation = smem_workspace.tensor(self.token_in_activation_smem_region, smem_base) - token_in_sf = smem_workspace.tensor(self.token_in_sf_smem_region, smem_base) warp_mbarrier = token_in_mbarriers + transfer_warp_idx warp_activation_stage = token_in_activation[transfer_warp_idx, None] - warp_sf_stage = token_in_sf[transfer_warp_idx, (None, None)] + if cutlass.const_expr(self.activation_is_quantized): + token_in_sf = smem_workspace.tensor(self.token_in_sf_smem_region, smem_base) + warp_sf_stage = token_in_sf[transfer_warp_idx, (None, None)] if lane_idx == Int32(0): cute.arch.mbarrier_init(warp_mbarrier, 1) cute.arch.sync_warp() fc1_activation_pointer = self._device_workspace.ptr(self.fc1_activation_region) - fc1_activation_sf = self.fc1_activation_sf_tensor(self._device_workspace) fc1_ready_counter = self._device_workspace.ptr(self.fc1_ready_region) activation_bytes = cute.cosize(warp_activation_stage) * int(self.activation_dtype.width) // 8 - activation_sf_bytes = cute.cosize(warp_sf_stage) * int(self.activation_sf_dtype.width) // 8 - sf_copy_elements = 4 - source_sf_values = cute.slice_(warp_sf_stage, (0, None)) - source_sf_vectors = cute.zipped_divide(source_sf_values, (sf_copy_elements,)) - sf_copy_atom = _copy_atom( - self.activation_sf_dtype, - sf_copy_elements * int(self.activation_sf_dtype.width), - ) + activation_sf_bytes = 0 + if cutlass.const_expr(self.activation_is_quantized): + fc1_activation_sf = self.fc1_activation_sf_tensor(self._device_workspace) + activation_sf_bytes = cute.cosize(warp_sf_stage) * int(self.activation_sf_dtype.width) // 8 + sf_copy_elements = 4 + source_sf_values = cute.slice_(warp_sf_stage, (0, None)) + source_sf_vectors = cute.zipped_divide(source_sf_values, (sf_copy_elements,)) + sf_copy_atom = _copy_atom(self.activation_sf_dtype, sf_copy_elements * int(self.activation_sf_dtype.width)) next_dense_token = global_warp_idx expert_valid_begin = Int32(0) - expert_sf_begin = Int32(0) + if cutlass.const_expr(self.activation_is_quantized): + expert_sf_begin = Int32(0) expert_ready_slot_begin = Int32(0) pull_phase = Int32(0) flag_tracker = make_flag_batch_tracker( @@ -2244,7 +2372,6 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No dense_token_idx = next_dense_token + Int32(pull_round) * global_warp_count token_in_expert = dense_token_idx - expert_valid_begin pool_token_idx = pool_expert_bases[local_expert] + token_in_expert - sf_token_idx = expert_sf_begin + token_in_expert ready_slot_idx = expert_ready_slot_begin + token_in_expert // Int32(self.tokens_per_fc1_ready_slot) metadata = TokenSrcMetadata.load( @@ -2256,77 +2383,110 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No + peer_offset + Int64(metadata.src_token) * Int64(self.bytes_per_token) ) - remote_sf = cute.make_tensor( - cute.make_ptr( - self._token_comm_args.activation_sf.dtype, - self._token_comm_args.activation_sf.iterator.toint() + peer_offset, - AddressSpace.gmem, - assumed_align=(self._token_comm_args.activation_sf.iterator.max_alignment), - ), - self._token_comm_args.activation_sf.layout, - ) - remote_sf_row = remote_sf[Int64(metadata.src_token), None] - - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - warp_mbarrier, Int32(activation_bytes + activation_sf_bytes) - ) - tma_load_1d( - warp_activation_stage.iterator, - Int64(remote_activation_address), - warp_mbarrier, - Int32(activation_bytes), - ) - tma_load_1d( - warp_sf_stage.iterator, - remote_sf_row.iterator, - warp_mbarrier, - Int32(activation_sf_bytes), + if cutlass.const_expr(self.activation_is_quantized): + remote_sf = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.activation_sf.dtype, + self._token_comm_args.activation_sf.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.activation_sf.iterator.max_alignment), + ), + self._token_comm_args.activation_sf.layout, ) - cute.arch.sync_warp() - cute.arch.mbarrier_wait(warp_mbarrier, pull_phase) + remote_sf_row = remote_sf[Int64(metadata.src_token), None] - destination_activation = cute.make_ptr( - self.activation_dtype, - fc1_activation_pointer.toint() + Int64(pool_token_idx) * Int64(self.bytes_per_token), - AddressSpace.gmem, - assumed_align=16, - ) - with cute.arch.elect_one(): - cp_async_bulk_s2g( - destination_activation, - warp_activation_stage.iterator, - Int32(activation_bytes), - ) - cute.arch.sync_warp() - cute.arch.cp_async_bulk_commit_group() - - destination_sf_row = fc1_activation_sf[Int64(sf_token_idx), ((None, None), None)] - destination_sf_values = cute.slice_(destination_sf_row, (0, None, None)) - destination_sf_values = cute.group_modes(destination_sf_values, 0, 2) - destination_sf_vectors = cute.zipped_divide(destination_sf_values, (sf_copy_elements,)) - sf_vector_count = cute.size(destination_sf_vectors, mode=[1]) - for sf_round in cutlass.range_constexpr(ceil_div(sf_vector_count, 32)): - sf_vector_idx = Int32(sf_round * 32) + lane_idx - if sf_vector_idx < Int32(sf_vector_count): - cute.copy( - sf_copy_atom, - source_sf_vectors[None, sf_vector_idx], - destination_sf_vectors[None, sf_vector_idx], + if cutlass.const_expr(self.token_in_stage_split): + # Chunked pull: the halved stage bounces each token row in two pieces, + # trading one extra NVLink round trip per token for SMEM. + for chunk_idx in cutlass.range_constexpr(self._token_in_stage_chunk_count): + chunk_offset = Int64(chunk_idx) * Int64(activation_bytes) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(warp_mbarrier, Int32(activation_bytes)) + tma_load_1d( + warp_activation_stage.iterator, + Int64(remote_activation_address) + chunk_offset, + warp_mbarrier, + Int32(activation_bytes), + ) + cute.arch.sync_warp() + cute.arch.mbarrier_wait(warp_mbarrier, pull_phase) + pull_phase = pull_phase ^ Int32(1) + + chunk_destination = cute.make_ptr( + self.activation_dtype, + fc1_activation_pointer.toint() + + Int64(pool_token_idx) * Int64(self.bytes_per_token) + + chunk_offset, + AddressSpace.gmem, + assumed_align=16, ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(chunk_destination, warp_activation_stage.iterator, Int32(activation_bytes)) + cute.arch.sync_warp() + cute.arch.cp_async_bulk_commit_group() + # Drain before the stage is reused by the next chunk (and before the flag). + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.sync_warp() + else: + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + warp_mbarrier, Int32(activation_bytes + activation_sf_bytes) + ) + tma_load_1d( + warp_activation_stage.iterator, + Int64(remote_activation_address), + warp_mbarrier, + Int32(activation_bytes), + ) + if cutlass.const_expr(self.activation_is_quantized): + tma_load_1d( + warp_sf_stage.iterator, remote_sf_row.iterator, warp_mbarrier, Int32(activation_sf_bytes) + ) + cute.arch.sync_warp() + cute.arch.mbarrier_wait(warp_mbarrier, pull_phase) - cute.arch.cp_async_bulk_wait_group(0, read=True) - cute.arch.sync_warp() + destination_activation = cute.make_ptr( + self.activation_dtype, + fc1_activation_pointer.toint() + Int64(pool_token_idx) * Int64(self.bytes_per_token), + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_activation, warp_activation_stage.iterator, Int32(activation_bytes)) + cute.arch.sync_warp() + cute.arch.cp_async_bulk_commit_group() + + if cutlass.const_expr(self.activation_is_quantized): + sf_token_idx = expert_sf_begin + token_in_expert + destination_sf_row = fc1_activation_sf[Int64(sf_token_idx), ((None, None), None)] + destination_sf_values = cute.slice_(destination_sf_row, (0, None, None)) + destination_sf_values = cute.group_modes(destination_sf_values, 0, 2) + destination_sf_vectors = cute.zipped_divide(destination_sf_values, (sf_copy_elements,)) + sf_vector_count = cute.size(destination_sf_vectors, mode=[1]) + for sf_round in cutlass.range_constexpr(ceil_div(sf_vector_count, 32)): + sf_vector_idx = Int32(sf_round * 32) + lane_idx + if sf_vector_idx < Int32(sf_vector_count): + cute.copy( + sf_copy_atom, + source_sf_vectors[None, sf_vector_idx], + destination_sf_vectors[None, sf_vector_idx], + ) + + if cutlass.const_expr(not self.token_in_stage_split): + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.sync_warp() ready_address = (fc1_ready_counter + ready_slot_idx).toint() flag_tracker = flag_tracker.accumulate(Int32(0), self.token_in_flag_batch, ready_address) cute.arch.sync_warp() - pull_phase = pull_phase ^ Int32(1) + if cutlass.const_expr(not self.token_in_stage_split): + pull_phase = pull_phase ^ Int32(1) next_dense_token = next_dense_token + pull_count * global_warp_count expert_valid_begin = expert_valid_end - expert_sf_begin = expert_sf_begin + ( - (expert_token_count + Int32(self.sf_padding_block - 1)) // Int32(self.sf_padding_block) - ) * Int32(self.sf_padding_block) + if cutlass.const_expr(self.activation_is_quantized): + expert_sf_begin = expert_sf_begin + ( + (expert_token_count + Int32(self.sf_padding_block - 1)) // Int32(self.sf_padding_block) + ) * Int32(self.sf_padding_block) expert_ready_slot_begin = expert_ready_slot_begin + ( (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) // Int32(self.tokens_per_fc1_ready_slot) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_deterministic.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_deterministic.py index af3e8b6545b2..6838edb87603 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_deterministic.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_deterministic.py @@ -41,6 +41,7 @@ _quant_spec = { + "bf16": (cutlass.BFloat16, None, 0), "nvfp4": (cutlass.Float4E2M1FN, cutlass.Float8E4M3FN, 16), "mxfp4": (cutlass.Float4E2M1FN, cutlass.Float8E8M0FNU, 32), "mxfp8_e4m3": (cutlass.Float8E4M3FN, cutlass.Float8E8M0FNU, 32), @@ -997,20 +998,42 @@ def _compute_push_tables(self, smem_base: cute.Pointer) -> None: iket.range_push("router.apply_receive_limit") if self._router_thread_idx == Int32(0): warp_totals[0] = Int32(0) + warp_totals[1] = Int32(0) cute.arch.sync_threads() - thread_local_total = Int32(0) - local_expert_rounds = ceil_div(self.experts_per_rank, block_thread_count) - for expert_round in cutlass.range_constexpr(local_expert_rounds): - local_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx - if local_expert < Int32(self.experts_per_rank): - thread_local_total = thread_local_total + sizes[owner_expert_begin + local_expert] - if thread_local_total > Int32(0): - cute.arch.atomic_add(warp_totals.iterator, thread_local_total, scope="cta") + + # sizes[] contains the same pre-truncation, all-source expert totals on + # every rank once raw_sizes_ready_target has been observed. Compute the + # group-wide policy bit locally so no post-kernel collective is needed. + # Slot 0 counts overflowing destinations; slot 1 records this rank's + # destination bit for local-only truncation. + rank_round_count = ceil_div(self.world_size, block_thread_count) + for rank_round in cutlass.range_constexpr(rank_round_count): + destination_rank = Int32(rank_round * block_thread_count) + self._router_thread_idx + if destination_rank < Int32(self.world_size): + destination_expert_begin = destination_rank * Int32(self.experts_per_rank) + raw_destination_total = Int32(0) + for local_expert in cutlass.range_constexpr(self.experts_per_rank): + raw_destination_total = ( + raw_destination_total + + sizes[destination_expert_begin + Int32(local_expert)] + ) + destination_overflow = Int32( + raw_destination_total > Int32(self.max_recv_size_per_rank) + ) + if destination_overflow: + cute.arch.atomic_add( + warp_totals.iterator, + Int32(1), + scope="cta", + ) + if destination_rank == Int32(self._router_local_rank): + warp_totals[1] = destination_overflow cute.arch.sync_threads() - raw_local_total = warp_totals[0] - did_overflow = Int32(raw_local_total > Int32(self.max_recv_size_per_rank)) + global_did_overflow = Int32(warp_totals[0] > Int32(0)) + local_did_overflow = warp_totals[1] if cutlass.const_expr(self.drop_on_overflow): + local_expert_rounds = ceil_div(self.experts_per_rank, block_thread_count) for expert_round in cutlass.range_constexpr(local_expert_rounds): local_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx if local_expert < Int32(self.experts_per_rank): @@ -1020,17 +1043,22 @@ def _compute_push_tables(self, smem_base: cute.Pointer) -> None: remaining_capacity = Int32(self.worst_case_token_count) - pool_expert_base[local_expert] if remaining_capacity > Int32(0): retained_size = cutlass.min(raw_size, remaining_capacity) - if did_overflow: + if local_did_overflow: sizes[owner_expert_begin + local_expert] = retained_size if self._router_thread_idx == Int32(0): - cute.arch.store(self._overflow_flag.iterator, did_overflow, sem="relaxed", scope="sys") + cute.arch.store( + self._overflow_flag.iterator, + global_did_overflow, + sem="relaxed", + scope="sys", + ) cute.arch.sync_threads() if cutlass.const_expr(self.drop_on_overflow): self._publish_push_table_and_sizes_ready(sizes_ready) else: - if did_overflow: + if global_did_overflow: if self._router_thread_idx == Int32(0): cute.arch.fence_acq_rel_sys() _device_trap() @@ -1391,6 +1419,8 @@ def impl_desc_require(cls) -> dict[str, type]: "token_back_schedule_mode": str, "reduce_topk_in_kernel": bool, "drop_on_overflow": bool, + "token_in_stage_split": OptionalRequirement(bool), + "external_dispatch_pool": OptionalRequirement(bool), } def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: @@ -1423,6 +1453,9 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.token_back_schedule_mode: TokenBackScheduleMode = impl_desc["token_back_schedule_mode"] self.reduce_topk_in_kernel = impl_desc["reduce_topk_in_kernel"] self.drop_on_overflow = impl_desc["drop_on_overflow"] + self.token_in_stage_split = impl_desc.get("token_in_stage_split", False) + self._token_in_stage_chunk_count = 2 if self.token_in_stage_split else 1 + self.external_dispatch_pool = impl_desc.get("external_dispatch_pool", False) self._validate_configuration() self._router = _MetadataPushRouter(problem_desc, impl_desc) @@ -1477,16 +1510,35 @@ def _validate_configuration(self) -> None: raise ValueError("fc2_done_signals_per_token_tile must be positive when token-back is enabled.") if self.reduce_topk_in_kernel and self.combine_format.act_dtype is not cutlass.BFloat16: raise ValueError("In-kernel top-k reduction requires BF16 combine data.") - element_block = self.activation_sf_vector_size * 4 - if self.hidden_size % element_block != 0: - raise ValueError(f"{self.quant_kind} requires hidden_size divisible by {element_block}.") + if self.activation_is_quantized: + element_block = self.activation_sf_vector_size * 4 + if self.hidden_size % element_block != 0: + raise ValueError(f"{self.quant_kind} requires hidden_size divisible by {element_block}.") + else: + # The 1D bulk copies move whole token rows, so each row must be whole 16-byte lines. + row_element_block = 128 // int(self.activation_dtype.width) + if self.hidden_size % row_element_block != 0: + raise ValueError(f"The BF16 wire requires hidden_size divisible by {row_element_block}.") if self.sf_padding_block % 128 != 0: raise ValueError("sf_padding_block must be a multiple of 128.") + if self.token_in_stage_split: + # Scale-free wire only: the SF pull and the token-back stage sizing both assume one + # whole-row stage per transfer warp. + if self.activation_is_quantized or self.token_back_enabled: + raise NotImplementedError( + "token_in_stage_split requires the BF16 wire with token_back disabled (epi_warps)." + ) + if self.bytes_per_token % (self._token_in_stage_chunk_count * 16) != 0: + raise ValueError("token_in_stage_split must cut a token row into whole 16-byte chunks.") @property def experts_per_rank(self) -> int: return self.expert_count // self.world_size + @property + def activation_is_quantized(self) -> bool: + return _quant_spec[self.quant_kind][1] is not None + @property def activation_dtype(self) -> type: return _quant_spec[self.quant_kind][0] @@ -1505,6 +1557,8 @@ def bytes_per_token(self) -> int: @property def activation_sf_hidden_padded(self) -> int: + if not self.activation_is_quantized: + return 0 valid_hidden = self.hidden_size // self.activation_sf_vector_size elements_per_16_bytes = 128 // int(self.activation_sf_dtype.width) return int(round_up(valid_hidden, elements_per_16_bytes)) @@ -1643,6 +1697,10 @@ def assign_device_members( local_rank: Int32, linear_cta_idx: Int32, ) -> None: + if cutlass.const_expr(self.external_dispatch_pool and token_comm_args.fc1_activation_pool is None): + raise ValueError("external_dispatch_pool requires TokenCommArgs.fc1_activation_pool.") + if cutlass.const_expr(not self.external_dispatch_pool and token_comm_args.fc1_activation_pool is not None): + raise ValueError("TokenCommArgs.fc1_activation_pool is only consumed under external_dispatch_pool.") self._device_workspace = device_workspace self._token_comm_args = token_comm_args self._local_rank = local_rank @@ -1678,22 +1736,24 @@ def _register_main_workspace(self, workspace: DeviceWorkspace) -> None: buffer_space="local", reset="tail_reset", ) - activation_element_count = self.worst_case_token_count * self.hidden_size - workspace.register( - self.fc1_activation_region, - self.activation_dtype, - (activation_element_count,), - buffer_space="local", - byte_alignment=128, - ) - activation_sf_element_count = self.worst_case_sf_token_count * self.activation_sf_hidden_padded - workspace.register( - self.fc1_activation_sf_region, - self.activation_sf_dtype, - (activation_sf_element_count,), - buffer_space="local", - byte_alignment=128, - ) + if not self.external_dispatch_pool: + activation_element_count = self.worst_case_token_count * self.hidden_size + workspace.register( + self.fc1_activation_region, + self.activation_dtype, + (activation_element_count,), + buffer_space="local", + byte_alignment=128, + ) + if self.activation_is_quantized: + activation_sf_element_count = self.worst_case_sf_token_count * self.activation_sf_hidden_padded + workspace.register( + self.fc1_activation_sf_region, + self.activation_sf_dtype, + (activation_sf_element_count,), + buffer_space="local", + byte_alignment=128, + ) if self.token_back_enabled: workspace.register( self.fc2_done_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local", reset="tail_reset" @@ -1751,16 +1811,17 @@ def register_smem_regions(self, workspace: SmemWorkspace) -> None: token_in_lifetime.register_tensor( self.token_in_activation_smem_region, self.activation_dtype, - (self.transfer_warp_count, self.hidden_size), - byte_alignment=16, - ) - token_in_lifetime.register_tensor( - self.token_in_sf_smem_region, - self.activation_sf_dtype, - (self.transfer_warp_count, (self.activation_sf_vector_size, self.activation_sf_hidden_padded)), - stride=(self.activation_sf_hidden_padded, (0, 1)), + (self.transfer_warp_count, self.hidden_size // self._token_in_stage_chunk_count), byte_alignment=16, ) + if self.activation_is_quantized: + token_in_lifetime.register_tensor( + self.token_in_sf_smem_region, + self.activation_sf_dtype, + (self.transfer_warp_count, (self.activation_sf_vector_size, self.activation_sf_hidden_padded)), + stride=(self.activation_sf_hidden_padded, (0, 1)), + byte_alignment=16, + ) if not self.token_back_enabled: return @@ -1775,7 +1836,11 @@ def register_smem_regions(self, workspace: SmemWorkspace) -> None: available_bytes_per_warp = self.standalone_chunk_bytes else: activation_bytes = self.bytes_per_token - sf_bytes = self.activation_sf_hidden_padded * int(self.activation_sf_dtype.width) // 8 + sf_bytes = ( + self.activation_sf_hidden_padded * int(self.activation_sf_dtype.width) // 8 + if self.activation_is_quantized + else 0 + ) available_bytes_per_warp = activation_bytes + sf_bytes if self.token_back_push_data: @@ -1811,15 +1876,24 @@ def register_smem_regions(self, workspace: SmemWorkspace) -> None: def fc1_ready_counter_pointer(self, device_workspace: DeviceWorkspace) -> cute.Pointer: return device_workspace.ptr(self.fc1_ready_region) + @property + def dispatch_pool_shape(self) -> Tuple[int, int]: + """Row-major ``(pool_capacity, hidden)`` geometry of the dispatch pool, workspace-owned or external.""" + return (self.worst_case_token_count, self.hidden_size) + @cute.jit def fc1_activation_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + if cutlass.const_expr(self.external_dispatch_pool): + raise ValueError("The dispatch pool is caller-owned under external_dispatch_pool; use the passed tensor.") return cute.make_tensor( device_workspace.ptr(self.fc1_activation_region), - cute.make_layout((self.worst_case_token_count, self.hidden_size), stride=(self.hidden_size, 1)), + cute.make_layout(self.dispatch_pool_shape, stride=(self.hidden_size, 1)), ) @cute.jit - def fc1_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + def fc1_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.activation_is_quantized): + return None layout = tile_atom_to_shape_SF( (self.worst_case_sf_token_count, self.hidden_size, 1), self.activation_sf_vector_size ) @@ -1913,7 +1987,10 @@ def _token_in_deduplicated(self, smem_workspace: SmemWorkspace, smem_base: cute. source_sf_vectors = cute.zipped_divide(source_sf_values, (4,)) sf_atom = _copy_atom(self.activation_sf_dtype, 4 * int(self.activation_sf_dtype.width)) destination_sf = self.fc1_activation_sf_tensor(self._device_workspace) - destination_data = self._device_workspace.ptr(self.fc1_activation_region) + if cutlass.const_expr(self.external_dispatch_pool): + destination_data = self._token_comm_args.fc1_activation_pool.iterator + else: + destination_data = self._device_workspace.ptr(self.fc1_activation_region) ready = self._device_workspace.ptr(self.fc1_ready_region) flags = make_flag_batch_tracker(use_async=self.token_in_async_flags, flag_address=Int64(0), accumulated_flags=Int32(0), phase=Int32(0), thread_idx=lane_idx) @@ -2110,27 +2187,34 @@ def token_in( iket.range_push("token_in.pull_payload") token_in_mbarriers = smem_workspace.ptr(self.token_in_mbarrier_region, smem_base) token_in_activation = smem_workspace.tensor(self.token_in_activation_smem_region, smem_base) - token_in_sf = smem_workspace.tensor(self.token_in_sf_smem_region, smem_base) warp_mbarrier = token_in_mbarriers + transfer_warp_idx warp_activation_stage = token_in_activation[transfer_warp_idx, None] - warp_sf_stage = token_in_sf[transfer_warp_idx, (None, None)] + if cutlass.const_expr(self.activation_is_quantized): + token_in_sf = smem_workspace.tensor(self.token_in_sf_smem_region, smem_base) + warp_sf_stage = token_in_sf[transfer_warp_idx, (None, None)] if lane_idx == Int32(0): cute.arch.mbarrier_init(warp_mbarrier, 1) cute.arch.sync_warp() - fc1_activation_pointer = self._device_workspace.ptr(self.fc1_activation_region) - fc1_activation_sf = self.fc1_activation_sf_tensor(self._device_workspace) + if cutlass.const_expr(self.external_dispatch_pool): + fc1_activation_pointer = self._token_comm_args.fc1_activation_pool.iterator + else: + fc1_activation_pointer = self._device_workspace.ptr(self.fc1_activation_region) fc1_ready_counter = self._device_workspace.ptr(self.fc1_ready_region) activation_bytes = cute.cosize(warp_activation_stage) * int(self.activation_dtype.width) // 8 - activation_sf_bytes = cute.cosize(warp_sf_stage) * int(self.activation_sf_dtype.width) // 8 - sf_copy_elements = 4 - source_sf_values = cute.slice_(warp_sf_stage, (0, None)) - source_sf_vectors = cute.zipped_divide(source_sf_values, (sf_copy_elements,)) - sf_copy_atom = _copy_atom(self.activation_sf_dtype, sf_copy_elements * int(self.activation_sf_dtype.width)) + activation_sf_bytes = 0 + if cutlass.const_expr(self.activation_is_quantized): + fc1_activation_sf = self.fc1_activation_sf_tensor(self._device_workspace) + activation_sf_bytes = cute.cosize(warp_sf_stage) * int(self.activation_sf_dtype.width) // 8 + sf_copy_elements = 4 + source_sf_values = cute.slice_(warp_sf_stage, (0, None)) + source_sf_vectors = cute.zipped_divide(source_sf_values, (sf_copy_elements,)) + sf_copy_atom = _copy_atom(self.activation_sf_dtype, sf_copy_elements * int(self.activation_sf_dtype.width)) next_dense_token = global_warp_idx expert_valid_begin = Int32(0) - expert_sf_begin = Int32(0) + if cutlass.const_expr(self.activation_is_quantized): + expert_sf_begin = Int32(0) expert_ready_slot_begin = Int32(0) pull_phase = Int32(0) flag_tracker = make_flag_batch_tracker( @@ -2154,7 +2238,6 @@ def token_in( dense_token_idx = next_dense_token + Int32(pull_round) * global_warp_count token_in_expert = dense_token_idx - expert_valid_begin pool_token_idx = pool_expert_base + token_in_expert - sf_token_idx = expert_sf_begin + token_in_expert ready_slot_idx = expert_ready_slot_begin + token_in_expert // Int32(self.tokens_per_fc1_ready_slot) metadata = TokenSrcMetadata.load( @@ -2166,73 +2249,112 @@ def token_in( + peer_offset + Int64(metadata.src_token) * Int64(self.bytes_per_token) ) - remote_sf = cute.make_tensor( - cute.make_ptr( - self._token_comm_args.activation_sf.dtype, - self._token_comm_args.activation_sf.iterator.toint() + peer_offset, - AddressSpace.gmem, - assumed_align=(self._token_comm_args.activation_sf.iterator.max_alignment), - ), - self._token_comm_args.activation_sf.layout, - ) - remote_sf_row = remote_sf[Int64(metadata.src_token), None] - - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - warp_mbarrier, Int32(activation_bytes + activation_sf_bytes) - ) - tma_load_1d( - warp_activation_stage.iterator, - Int64(remote_activation_address), - warp_mbarrier, - Int32(activation_bytes), - ) - tma_load_1d( - warp_sf_stage.iterator, remote_sf_row.iterator, warp_mbarrier, Int32(activation_sf_bytes) + if cutlass.const_expr(self.activation_is_quantized): + remote_sf = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.activation_sf.dtype, + self._token_comm_args.activation_sf.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.activation_sf.iterator.max_alignment), + ), + self._token_comm_args.activation_sf.layout, ) - cute.arch.sync_warp() - cute.arch.mbarrier_wait(warp_mbarrier, pull_phase) - - destination_activation = cute.make_ptr( - self.activation_dtype, - fc1_activation_pointer.toint() + Int64(pool_token_idx) * Int64(self.bytes_per_token), - AddressSpace.gmem, - assumed_align=16, - ) - with cute.arch.elect_one(): - cp_async_bulk_s2g(destination_activation, warp_activation_stage.iterator, Int32(activation_bytes)) - cute.arch.sync_warp() - cute.arch.cp_async_bulk_commit_group() + remote_sf_row = remote_sf[Int64(metadata.src_token), None] - destination_sf_row = fc1_activation_sf[Int64(sf_token_idx), ((None, None), None)] - destination_sf_values = cute.slice_(destination_sf_row, (0, None, None)) - destination_sf_values = cute.group_modes(destination_sf_values, 0, 2) - destination_sf_vectors = cute.zipped_divide(destination_sf_values, (sf_copy_elements,)) - sf_vector_count = cute.size(destination_sf_vectors, mode=[1]) - for sf_round in cutlass.range_constexpr(ceil_div(sf_vector_count, 32)): - sf_vector_idx = Int32(sf_round * 32) + lane_idx - if sf_vector_idx < Int32(sf_vector_count): - cute.copy( - sf_copy_atom, - source_sf_vectors[None, sf_vector_idx], - destination_sf_vectors[None, sf_vector_idx], + if cutlass.const_expr(self.token_in_stage_split): + # Chunked pull: the halved stage bounces each token row in two pieces, + # trading one extra NVLink round trip per token for SMEM. + for chunk_idx in cutlass.range_constexpr(self._token_in_stage_chunk_count): + chunk_offset = Int64(chunk_idx) * Int64(activation_bytes) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(warp_mbarrier, Int32(activation_bytes)) + tma_load_1d( + warp_activation_stage.iterator, + Int64(remote_activation_address) + chunk_offset, + warp_mbarrier, + Int32(activation_bytes), + ) + cute.arch.sync_warp() + cute.arch.mbarrier_wait(warp_mbarrier, pull_phase) + pull_phase = pull_phase ^ Int32(1) + + chunk_destination = cute.make_ptr( + self.activation_dtype, + fc1_activation_pointer.toint() + + Int64(pool_token_idx) * Int64(self.bytes_per_token) + + chunk_offset, + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(chunk_destination, warp_activation_stage.iterator, Int32(activation_bytes)) + cute.arch.sync_warp() + cute.arch.cp_async_bulk_commit_group() + # Drain before the stage is reused by the next chunk (and before the flag). + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.sync_warp() + else: + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + warp_mbarrier, Int32(activation_bytes + activation_sf_bytes) ) + tma_load_1d( + warp_activation_stage.iterator, + Int64(remote_activation_address), + warp_mbarrier, + Int32(activation_bytes), + ) + if cutlass.const_expr(self.activation_is_quantized): + tma_load_1d( + warp_sf_stage.iterator, remote_sf_row.iterator, warp_mbarrier, Int32(activation_sf_bytes) + ) + cute.arch.sync_warp() + cute.arch.mbarrier_wait(warp_mbarrier, pull_phase) + + destination_activation = cute.make_ptr( + self.activation_dtype, + fc1_activation_pointer.toint() + Int64(pool_token_idx) * Int64(self.bytes_per_token), + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_activation, warp_activation_stage.iterator, Int32(activation_bytes)) + cute.arch.sync_warp() + cute.arch.cp_async_bulk_commit_group() + + if cutlass.const_expr(self.activation_is_quantized): + sf_token_idx = expert_sf_begin + token_in_expert + destination_sf_row = fc1_activation_sf[Int64(sf_token_idx), ((None, None), None)] + destination_sf_values = cute.slice_(destination_sf_row, (0, None, None)) + destination_sf_values = cute.group_modes(destination_sf_values, 0, 2) + destination_sf_vectors = cute.zipped_divide(destination_sf_values, (sf_copy_elements,)) + sf_vector_count = cute.size(destination_sf_vectors, mode=[1]) + for sf_round in cutlass.range_constexpr(ceil_div(sf_vector_count, 32)): + sf_vector_idx = Int32(sf_round * 32) + lane_idx + if sf_vector_idx < Int32(sf_vector_count): + cute.copy( + sf_copy_atom, + source_sf_vectors[None, sf_vector_idx], + destination_sf_vectors[None, sf_vector_idx], + ) if cutlass.const_expr(self.token_in_flag_batch == 1): cute.arch.cp_async_bulk_wait_group(0) - else: + elif cutlass.const_expr(not self.token_in_stage_split): cute.arch.cp_async_bulk_wait_group(0, read=True) cute.arch.sync_warp() ready_address = (fc1_ready_counter + ready_slot_idx).toint() flag_tracker = flag_tracker.accumulate(Int32(0), self.token_in_flag_batch, ready_address) cute.arch.sync_warp() - pull_phase = pull_phase ^ Int32(1) + if cutlass.const_expr(not self.token_in_stage_split): + pull_phase = pull_phase ^ Int32(1) next_dense_token = next_dense_token + pull_count * global_warp_count expert_valid_begin = expert_valid_end - expert_sf_begin = expert_sf_begin + ( - (expert_token_count + Int32(self.sf_padding_block - 1)) // Int32(self.sf_padding_block) - ) * Int32(self.sf_padding_block) + if cutlass.const_expr(self.activation_is_quantized): + expert_sf_begin = expert_sf_begin + ( + (expert_token_count + Int32(self.sf_padding_block - 1)) // Int32(self.sf_padding_block) + ) * Int32(self.sf_padding_block) expert_ready_slot_begin = expert_ready_slot_begin + ( (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) // Int32(self.tokens_per_fc1_ready_slot) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/megamoe_aot.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/megamoe_aot.py index a2e46b2ecbf5..51534b776661 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/megamoe_aot.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/megamoe_aot.py @@ -10,5 +10,9 @@ def make_megamoe_aot_callable(function, kernel): argument_names = list(signature(kernel.__call__).parameters) # Home-only entries can omit the unused READY pair, as with direct JIT # calls. Helper-enabled entries require both READY operands explicitly. - defaults = (None, None) if kernel.helper_expert_count == 0 else () + # Kernels that take the trailing optional fc1_activation_per_token_scale + # may omit it too (defaults are right-aligned, hence its position last). + trailing_optional = 1 if argument_names[-1] == "fc1_activation_per_token_scale" else 0 + ready_optional = 2 if kernel.helper_expert_count == 0 else 0 + defaults = (None,) * (ready_optional + trailing_optional) return make_kwargs_wrapper(function, argument_names, arg_defaults=defaults) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py index 8057166cabff..a363fcdd831a 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py @@ -32,7 +32,7 @@ from .....helpers.constants import Nvfp4E2M1RcpLimit, Fp8E4M3RcpLimit, Fp8E5M2RcpLimit, Fp32Max from .....helpers.cute_py_helpers import tcgen05_block_scaled_acc_dtype from .....helpers.dsl_helpers import mark_alignment -from .....helpers.flag_batch import GpuReleaseFlagBatchTracker +from .....helpers.flag_batch import make_flag_batch_tracker from .....helpers.ptx_helpers import ( cp_async_bulk_s2g, cp_reduce_async_bulk_add_bf16_s2g, @@ -915,6 +915,11 @@ class GatedActEpilogueArgs: # ----------------------------------- # MoE domain (token, topk), deepgemm graph only? for transformer graph, we want reduce kernel to perform the score mul. topk_scores: Optional[cute.Tensor] + # ----------------------------------- + # Pool-ordered (token,) fp32. The FC1 activation was quantized from ``x / s_t``, so the FC1 + # epilogue multiplies ``s_t`` back onto the raw accumulator before clamp / activation. It has + # no other home: the e4m3 block scale cannot carry an fp32 factor and alpha is per-expert. + fc1_activation_per_token_scale: Optional[cute.Tensor] = None class SwapABGatedActEpilogue(KernelComponent): @@ -944,6 +949,10 @@ class SwapABGatedActEpilogue(KernelComponent): # warp 32 accumulator rows, which the gate/up interleave halves. _EpilogueFc1IntermediateDownPerWarp = 16 smem_scratch_overlay: ClassVar[str] = "blackwell.swap_ab_gated_act_epilogue.scratch" + # (subtile, warp, 32) fp32 de-interleave buffer for the per-token FC1 activation scale (see + # SwapABFc1Epilogue.run_subtile): 512 B per subtile, 2 KB at cta_tile_n=256. Registered only + # when the descriptor enables the scale (``fc1_activation_per_token_scale``). + fc1_act_scale_scratch_region: ClassVar[str] = "blackwell.swap_ab_gated_act_epilogue.fc1_act_scale_scratch" fc1_staging_region: ClassVar[str] = "blackwell.swap_ab_gated_act_epilogue.fc1_staging" fc2_staging_region: ClassVar[str] = "blackwell.swap_ab_gated_act_epilogue.fc2_staging" # A 144 B FP8 stride rotates each row by four SMEM banks. @@ -965,8 +974,17 @@ def problem_desc_require(cls) -> dict[str, type]: # existing SwiGLU descriptors stay valid unchanged. See _resolve_situ_betas. "situ_beta": OptionalRequirement(Optional[float]), "situ_linear_beta": OptionalRequirement(Optional[float]), + # "swiglu" (default) or "quick_geglu": up * gate * sigmoid(1.702 * gate). Same core as + # SwiGLU with the sigmoid argument scaled, so it costs nothing extra; exclusive with SiTU. + "gate_activation": OptionalRequirement(str), + # Whether run() will receive the per-token FC1 activation scale. Absent means False: + # the SMEM scratch it needs is reserved only for descriptors that enable it. + "fc1_activation_per_token_scale": OptionalRequirement(bool), } + _supported_gate_activations: ClassVar[Tuple[str, ...]] = ("swiglu", "quick_geglu") + _quick_geglu_sigmoid_scale: ClassVar[float] = 1.702 + @classmethod def impl_desc_require(cls) -> dict[str, object]: return { @@ -979,6 +997,7 @@ def impl_desc_require(cls) -> dict[str, object]: "fc2_epi_flag_batch": int, "fc2_tma_stages": OptionalRequirement(int), "reduce_topk_in_kernel": OptionalRequirement(bool), + "epi_async_release": OptionalRequirement(bool), "token_back_push_data": OptionalRequirement(bool), } @@ -1005,6 +1024,27 @@ def _resolve_situ_betas(problem_desc: ProblemDesc) -> Tuple[Optional[float], Opt raise ValueError("SiTU does not support SwiGLUBias constants.") return float(beta), float(linear_beta) + @classmethod + def _resolve_gate_activation(cls, problem_desc: ProblemDesc) -> str: + """Resolve the gated-activation core name; absent means SwiGLU. + + quick_geglu only rescales the sigmoid argument of the SwiGLU core (1.702), so it keeps the + clamp semantics; SiTU is a different core and cannot be combined with it. Baked in at + codegen time, so the enclosing KernelClass folds it into ``name()``. + """ + gate_activation = problem_desc.get("gate_activation") + if gate_activation is None: + return "swiglu" + if gate_activation not in cls._supported_gate_activations: + raise ValueError( + f"gate_activation must be one of {cls._supported_gate_activations}, got {gate_activation!r}." + ) + if gate_activation != "swiglu" and problem_desc.get("situ_beta") is not None: + raise ValueError(f"gate_activation={gate_activation!r} cannot be combined with SiTU betas.") + if gate_activation != "swiglu" and problem_desc.get("swiglu_alpha") is not None: + raise ValueError(f"gate_activation={gate_activation!r} cannot be combined with SwiGLUBias constants.") + return gate_activation + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self._validate_desc_inputs(problem_desc, impl_desc) @@ -1022,10 +1062,16 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.swiglu_alpha = float(self.swiglu_alpha) self.swiglu_beta = float(self.swiglu_beta) self.situ_beta, self.situ_linear_beta = self._resolve_situ_betas(problem_desc) + self.gate_activation = self._resolve_gate_activation(problem_desc) + # The scratch costs mainloop SMEM budget, so only a descriptor that enables the scale + # reserves it; passing the tensor without the flag is rejected in run(). + self.act_scale_scratch_enabled = bool(problem_desc.get("fc1_activation_per_token_scale", False)) self.mma_tiler_mnk = impl_desc["mma_tiler_mnk"] self.cluster_shape_mn = impl_desc["cluster_shape_mn"] self.use_2cta_instrs = impl_desc["use_2cta_instrs"] self.fc2_use_bulk = impl_desc["fc2_use_bulk"] + # fc1/fc2 done flags published with red.async.release (default) instead of the blocking red.release (see flag_batch). + self.epi_async_release = impl_desc.get("epi_async_release", True) self.communication_enabled = impl_desc["communication_enabled"] self.fc1_epi_flag_batch = impl_desc["fc1_epi_flag_batch"] self.fc2_epi_flag_batch = impl_desc["fc2_epi_flag_batch"] @@ -1057,6 +1103,9 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.reduce_topk_in_epilogue = self.reduce_topk_in_kernel and not self.token_back_push_data if not 1 <= self.fc1_epi_flag_batch <= 32 or not 1 <= self.fc2_epi_flag_batch <= 32: raise ValueError("Epilogue flag batch sizes must be in [1, 32].") + if self.epi_async_release and max(self.fc1_epi_flag_batch, self.fc2_epi_flag_batch) > self._EpilogueWarpCnt: + # The async tracker batches per warp (one elected publisher per warp), the blocking one per lane. + raise ValueError(f"epi_async_release=True needs epilogue flag batch sizes <= {self._EpilogueWarpCnt}.") self.cluster_tile_intermediate_downproj = self._EpilogueFc1IntermediateDownTileSize * self.cluster_shape_mn[0] atom_thr_size = 2 if self.use_2cta_instrs else 1 @@ -1139,6 +1188,15 @@ def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: FC1's staging tile and FC2's store tile belong to different work tiles, separated by the tile-boundary TMA drain and rendezvous in ``run()``, so they never coexist. """ + if self.act_scale_scratch_enabled: + # (subtile, warp, 32): every subtile of the work tile keeps its own slot so the whole + # tile's scales can be staged once, ahead of the accumulator wait. + smem_workspace.register_tensor( + self.fc1_act_scale_scratch_region, + cutlass.Float32, + (self.subtile_cnt * self._EpilogueWarpCnt * 32,), + byte_alignment=64, + ) overlay = smem_workspace.create_overlay(self.smem_scratch_overlay) overlay.add_lifetime("fc1_staging").register_tensor( self.fc1_staging_region, cutlass.Int8, (self.fc1_staging_bytes,), byte_alignment=128 @@ -1298,6 +1356,15 @@ def run( fc1_alpha=None, fc2_alpha=None, fc1_norm_const=None, topk_scores=None ) fc1_staging_pointer = smem_workspace.ptr(self.fc1_staging_region, smem_base) + if cutlass.const_expr(self.act_scale_scratch_enabled): + fc1_act_scale_scratch_pointer = smem_workspace.ptr(self.fc1_act_scale_scratch_region, smem_base) + else: + if cutlass.const_expr(optional_epi_args.fc1_activation_per_token_scale is not None): + raise ValueError( + "fc1_activation_per_token_scale was passed, but the descriptor does not enable it " + "(problem_desc['fc1_activation_per_token_scale'] must be True), so no SMEM scratch was reserved." + ) + fc1_act_scale_scratch_pointer = None if cutlass.const_expr(self.fc2_tma_stages > 0): fc2_smem_tensor = smem_workspace.tensor(self.fc2_staging_region, smem_base) else: @@ -1321,6 +1388,7 @@ def run( self, tidx, fc1_staging_pointer, + fc1_act_scale_scratch_pointer, kernel_extension, tma_atom_fc1_output, fc1_output, @@ -1348,8 +1416,20 @@ def run( wait_only_named_barrier = self.epilogue_sync_barrier() is_odd_turn = cutlass.Int32(1) work_tile_info = sched_consumer.consume_work() - - flag_tracker = GpuReleaseFlagBatchTracker( + # With per-token scales the epilogue runs one work tile ahead on the scheduler ring: the + # next tile's scale loads are issued before the current tile is processed, so their + # latency hides behind a whole tile instead of sitting in front of the activation. A + # prefetch may only be issued once at least one acc wait lies behind it: in the EP kernel + # that wait is what orders the pool reads after every rank's router pushes. So the CTA's + # first FC1 tile loads its own scales after its acc wait (the one exposed load per CTA) and + # prefetches the next tile from there; every other tile is prefetched at the loop top. + prefetch_act_scales = cutlass.const_expr(optional_epi_args.fc1_activation_per_token_scale is not None) + # Whether the CURRENT tile's first-half scales sit in the prefetch registers. + act_scales_prefetched = cutlass.Int32(0) + loop_iteration = cutlass.Int32(0) + + flag_tracker = make_flag_batch_tracker( + self.epi_async_release, flag_address=Int64(0), accumulated_flags=cutlass.Int32(0), phase=cutlass.Int32(work_tile_info.phase), @@ -1357,6 +1437,23 @@ def run( ) while work_tile_info.is_valid_tile: + if cutlass.const_expr(prefetch_act_scales): + # The current tile is valid, so the scheduler publishes at least a terminal tile + # after it: consuming here cannot block forever. Stage the current tile's scales + # (from the registers prefetched last iteration) before those registers are + # reloaded for the next tile. + next_work_tile_info = sched_consumer.consume_work() + next_act_scales_prefetched = cutlass.Int32(0) + if work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1): + if act_scales_prefetched == cutlass.Int32(1): + fc1_epi.stage_act_scales(work_tile_info) + fc1_epi.prefetch_act_scales(next_work_tile_info) + # else: fc1_epi loads + stages this tile after its acc wait, then prefetches + # the next one itself (the registers are busy until then). + next_act_scales_prefetched = cutlass.Int32(1) + elif loop_iteration > cutlass.Int32(0): + fc1_epi.prefetch_act_scales(next_work_tile_info) + next_act_scales_prefetched = cutlass.Int32(1) if cutlass.const_expr(self.overlapping_accum): tmem_stage_idx = acc_consumer_state.phase else: @@ -1370,6 +1467,8 @@ def run( acc_pipeline=acc_pipeline, acc_consumer_state=acc_consumer_state, is_odd_turn=is_odd_turn, + act_scales_staged=act_scales_prefetched, + next_work_tile_info=next_work_tile_info if cutlass.const_expr(prefetch_act_scales) else None, ) else: # The __call__ args should only take the while loop args, leave all loop irrevalent args to the init. @@ -1380,6 +1479,9 @@ def run( acc_consumer_state=acc_consumer_state, is_odd_turn=is_odd_turn, ) + if cutlass.const_expr(prefetch_act_scales): + act_scales_prefetched = next_act_scales_prefetched + loop_iteration = loop_iteration + cutlass.Int32(1) prev_work_tile_info = work_tile_info cur_was_linear1 = prev_work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1) @@ -1387,7 +1489,10 @@ def run( if cutlass.const_expr(self.overlapping_accum): is_odd_turn = cutlass.Int32(1) - is_odd_turn - work_tile_info = sched_consumer.consume_work() + if cutlass.const_expr(prefetch_act_scales): + work_tile_info = next_work_tile_info + else: + work_tile_info = sched_consumer.consume_work() # Every asynchronous store commits at issue; drain before completion or scratch reuse. cute.arch.cp_async_bulk_wait_group(0) @@ -1422,6 +1527,7 @@ def __init__( base: SwapABGatedActEpilogue, tidx: cutlass.Int32, staging_pointer: cute.Pointer, + act_scale_scratch_pointer: cute.Pointer, kernel_extension: BlockScaledSwapAbFc12Extension, tma_atom_fc1_output: cute.CopyAtom, fc1_output: cute.Tensor, # fake (m,n,l) domain @@ -1442,6 +1548,22 @@ def __init__( # Kept unswizzled and untyped so fc1_quant can carve an fp32 amax-exchange view out of one # staging stage without going through the staging tensor's swizzle. self.staging_pointer = staging_pointer + # (subtile_cnt * warp_cnt * 32,) fp32: per (subtile, warp) slot holding the first-half + # token scales de-interleaved by token parity (see stage_act_scales / run_subtile). + # Prefetch registers hold the NEXT FC1 tile's first-half scales (one per subtile, this + # lane's token) and its expert alpha: loaded while the current tile is processed so the + # global-load latency never sits in front of the activation. + if optional_epi_args.fc1_activation_per_token_scale is not None: + self.act_scale_scratch = cute.make_tensor( + cute.recast_ptr(act_scale_scratch_pointer, dtype=cutlass.Float32), + cute.make_layout((base.subtile_cnt * base._EpilogueWarpCnt * 32,)), + ) + self.prefetched_act_scales = cute.make_rmem_tensor((base.subtile_cnt,), cutlass.Float32) + self.prefetched_act_alpha = cute.make_rmem_tensor((1,), cutlass.Float32) + else: + self.act_scale_scratch = None + self.prefetched_act_scales = None + self.prefetched_act_alpha = None self.kernel_extension = kernel_extension self.fc1_tma_atom = tma_atom_fc1_output self.fc1_output = fc1_output @@ -1465,6 +1587,57 @@ def __new_from_mlir_values__(self, values: List[ir.Value]) -> "SwapABFc1Epilogue assert len(values) == 0 return self + @cute.jit + def prefetch_act_scales(self, work_tile_info: SwapAbFc12WorkTileInfo) -> None: + """Issue the global loads for ``work_tile_info``'s first-half token scales (and its alpha). + + Normally called one tile ahead of the tile's processing, so the loads land while the + current tile runs; a CTA's first FC1 tile calls it for itself after its acc wait. Only FC1 + tiles load, and only lanes whose token is valid (the pool view ends at the expert's valid + rows). The values stay in this lane's registers until stage_act_scales moves them to SMEM. + """ + if cutlass.const_expr(self.optional_epi_args.fc1_activation_per_token_scale is None): + return + if work_tile_info.is_valid_tile: + if work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1): + act_scale_tensor, _ = self.kernel_extension.get_gmem_tensor( + "fc1_act_scale", self.optional_epi_args.fc1_activation_per_token_scale, work_tile_info + ) + if cutlass.const_expr(self.optional_epi_args.fc1_alpha is not None): + self.prefetched_act_alpha[0] = self.optional_epi_args.fc1_alpha[work_tile_info.expert_idx] + tile_token_base = work_tile_info.tile_n_idx * self.cta_tile_n + valid_tokens = work_tile_info.valid_tokens_in_cta_tile + for subtile in cutlass.range_constexpr(self.subtile_cnt): + token_in_tile = cutlass.Int32(subtile * self._EpilogueTokenTileSize) + self.lane_idx + if token_in_tile < valid_tokens: + self.prefetched_act_scales[subtile] = act_scale_tensor[tile_token_base + token_in_tile] + + @cute.jit + def stage_act_scales(self, work_tile_info: SwapAbFc12WorkTileInfo) -> None: + """Move the prefetched first-half scales of ``work_tile_info`` into this warp's SMEM slots. + + Per subtile, lane L owns token L of the first 32 tokens; it stores the alpha-folded scale + at slot (L%2)*16 + L//2 of the (subtile, warp) scratch row, so each token parity's 16 + scales are contiguous for run_subtile's four 16-byte loads. Lanes past the valid tokens + hold no scale (nothing was loaded for them) and stage 1.0, like the second half's default, + so the pad rows the subtile store still writes stay deterministic. + """ + if cutlass.const_expr(self.optional_epi_args.fc1_activation_per_token_scale is None): + return + valid_tokens = work_tile_info.valid_tokens_in_cta_tile + own_slot_offset = (self.lane_idx % cutlass.Int32(2)) * cutlass.Int32(16) + self.lane_idx // 2 + for subtile in cutlass.range_constexpr(self.subtile_cnt): + staged_scale = cutlass.Float32(1.0) + if cutlass.Int32(subtile * self._EpilogueTokenTileSize) + self.lane_idx < valid_tokens: + folded_scale = self.prefetched_act_scales[subtile] + if cutlass.const_expr(self.optional_epi_args.fc1_alpha is not None): + folded_scale = folded_scale * self.prefetched_act_alpha[0] + staged_scale = folded_scale + self.act_scale_scratch[ + cutlass.Int32((subtile * self._EpilogueWarpCnt) * 32) + self.warp_idx * cutlass.Int32(32) + own_slot_offset + ] = staged_scale + cute.arch.sync_warp() + @cute.jit def signal_fc1_done(self, work_tile_info, next_work_tile_info, flag_tracker): # Only in-bound intermediate_downproj tiles signal; OOB -> null slot. @@ -1490,6 +1663,11 @@ def __call__( acc_pipeline, acc_consumer_state, is_odd_turn: cutlass.Int32, + # 0 for a CTA's first FC1 tile: nothing was prefetched for it, so after the acc wait it + # loads and stages its own scales and then prefetches ``next_work_tile_info``'s (see + # SwapABGatedActEpilogue.run). + act_scales_staged: cutlass.Int32, + next_work_tile_info: Optional[SwapAbFc12WorkTileInfo] = None, ): # (tokens_this_expert, intermediate_down, 1) real_fc1_output, _ = self.kernel_extension.get_gmem_tensor("c", self.fc1_output, work_tile_info) @@ -1504,13 +1682,29 @@ def __call__( norm_const = self.optional_epi_args.fc1_norm_const[work_tile_info.expert_idx] else: norm_const = None + valid_tokens = work_tile_info.valid_tokens_in_cta_tile + if cutlass.const_expr(self.optional_epi_args.fc1_activation_per_token_scale is not None): + # (tokens_this_expert,), same pool-ordered view as topk scores. The first-half scales + # were staged into SMEM by stage_act_scales before this call. + act_scale_tensor, _ = self.kernel_extension.get_gmem_tensor( + "fc1_act_scale", self.optional_epi_args.fc1_activation_per_token_scale, work_tile_info + ) + else: + act_scale_tensor = None # (cta_tile_m, cta_tile_n) -> (epi_tile_m, epi_tile_n, iters) tmem_acc_tensor_tiled_by_epi_tile = cute.flat_divide( tmem_acc_tensor, (self._EpilogueFc1IntermediateGateUpTileSize, self._EpilogueTokenTileSize) )[None, None, 0, None] acc_pipeline.consumer_wait(acc_consumer_state) - valid_tokens = work_tile_info.valid_tokens_in_cta_tile + if cutlass.const_expr(act_scale_tensor is not None): + # The acc wait orders these pool reads after the router's writes (dispatch pulled these + # tokens only after metadata_ready). Exposed once per CTA: the next tile's prefetch is + # issued right here, so every later FC1 tile finds its scales already in registers. + if act_scales_staged == cutlass.Int32(0): + self.prefetch_act_scales(work_tile_info) + self.stage_act_scales(work_tile_info) + self.prefetch_act_scales(next_work_tile_info) # Overlap path preloads two subtiles before releasing acc TMEM. unroll_tile_cnt = 2 if cutlass.const_expr(self.overlapping_accum) else 0 @@ -1563,6 +1757,7 @@ def __call__( fc1_output_sf=real_fc1_output_sf, alpha_val=alpha_val, norm_const=norm_const, + act_scale_tensor=act_scale_tensor, ) for i in cutlass.range(remain_subtile_cnt, unroll=1): @@ -1582,6 +1777,7 @@ def __call__( fc1_output_sf=real_fc1_output_sf, alpha_val=alpha_val, norm_const=norm_const, + act_scale_tensor=act_scale_tensor, ) # Non-overlap-path release: at the natural task-tile boundary. @@ -1603,6 +1799,8 @@ def run_subtile( fc1_output_sf: cute.Tensor, alpha_val: Optional[cutlass.Float32], norm_const: Optional[cutlass.Float32], + # (tokens_this_expert,) fp32 per-token activation scale view, or None. + act_scale_tensor: Optional[cute.Tensor] = None, ): if cutlass.const_expr(self.optional_epi_args.topk_scores is not None): # This means we need to perform DeepGEMM computation graph, topk_score at fc1 pre-quant @@ -1633,6 +1831,53 @@ def run_subtile( ) else: topk_scores = None + # Per-token activation scale, with the expert alpha folded in once per lane so the + # activation's dequant multiply stays a single packed FMUL per register pair. + # second half (token 32..63): activated after its transpose, one token per lane -> the + # lane's own scalar. + # first half (token 0..31): activated before its transpose, where register i of lane L + # belongs to token 2i + (L//2)%2 -> the lane needs the 16 scales of one parity. Lane L + # writes token L's scale to slot (L%2)*16 + L//2 of this warp's scratch (even tokens + # first, then odd), so each parity's 16 scales are contiguous and come back with four + # 16-byte LDS instead of 16 shuffles. + # Both halves end up as a complete dequant multiplier (alpha x s_t), so alpha_swiglu_clamp + # receives it INSTEAD of alpha_val and multiplies exactly once per register. + if cutlass.const_expr(act_scale_tensor is not None): + # Second half: one token per lane (token lane+32). Loaded here, alpha folded here; it is + # consumed only after the 32x32 transpose, which hides the load latency. Lanes past the + # valid tokens skip the load (the pool view ends at the expert's valid rows). + dequant_scale_token_32_64 = cutlass.Float32(1.0) + second_half_token_in_tile = subtile_idx * cutlass.Int32(self._EpilogueTokenTileSize) + self.lane_idx + 32 + if second_half_token_in_tile < work_tile_info.valid_tokens_in_cta_tile: + dequant_scale_token_32_64 = act_scale_tensor[current_two_token_idices[1]] + if cutlass.const_expr(alpha_val is not None): + dequant_scale_token_32_64 = dequant_scale_token_32_64 * alpha_val + # First half: 16 tokens per lane. The 16 values were loaded (prefetch_act_scales), + # multiplied by alpha and written de-interleaved to SMEM by stage_act_scales before this + # tile started; here they are only read back. Alpha is therefore already included. + lane_parity = (self.lane_idx // cutlass.Int32(2)) % cutlass.Int32(2) + scale_vec_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=128) + # ((subtile*warps + warp)*32 + parity*16) * 4 B is a multiple of 64 B; the dynamic + # offset hides that from the copy verifier. + parity_scales_smem = mark_alignment( + cute.make_tensor( + self.act_scale_scratch.iterator + + ( + (subtile_idx * cutlass.Int32(self._EpilogueWarpCnt) + self.warp_idx) * cutlass.Int32(32) + + lane_parity * cutlass.Int32(16) + ), + cute.make_layout((16,)), + ), + 64, + ) + dequant_scales_token_0_32 = cute.make_rmem_tensor((16,), cutlass.Float32) + cute.copy(scale_vec_atom, parity_scales_smem, dequant_scales_token_0_32) + # With a per-token scale the dequant multiplier carries alpha; do not apply it twice. + activation_alpha = None + else: + dequant_scales_token_0_32 = None + dequant_scale_token_32_64 = None + activation_alpha = alpha_val # Step 0: load tmem if cutlass.const_expr(preload_acc is not None): @@ -1673,7 +1918,9 @@ def run_subtile( ) # Step 1: perform swiglu on the first part, interleave with the second's 32x32 tmem transpose. - token_0_32_pre_quant_pre_trans = self.alpha_swiglu_clamp(gate_token_0_32, up_token_0_32, alpha_val) + token_0_32_pre_quant_pre_trans = self.alpha_swiglu_clamp( + gate_token_0_32, up_token_0_32, activation_alpha, dequant_scales=dequant_scales_token_0_32 + ) # gate_token_32_64 / up_token_32_64 are already in the transpose input # distribution (see TmemTranspose16x32 / load_subtile_raw_acc). @@ -1688,7 +1935,10 @@ def run_subtile( ) token_32_64_pre_quant = self.alpha_swiglu_clamp( - gate_token_32_64_trans_pre_act, up_token_32_64_trans_pre_act, alpha_val + gate_token_32_64_trans_pre_act, + up_token_32_64_trans_pre_act, + activation_alpha, + dequant_scale=dequant_scale_token_32_64, ) token_0_32_tmem_trans = TmemTranspose16x32( @@ -1746,6 +1996,13 @@ def alpha_swiglu_clamp( gate_rmem: cute.Tensor, # Raw fc1 acc (pre-dequant); even-size 1D fp32 rmem up_rmem: cute.Tensor, # Raw fc1 acc (pre-dequant); even-size 1D fp32 rmem alpha_val: Optional[cutlass.Float32], + *, + # Full dequant multiplier with alpha ALREADY folded in (the caller multiplies once per + # lane), replacing alpha_val when given. Exactly one form applies: + # dequant_scale : one token per lane (post-transpose distribution), one scalar. + # dequant_scales: (n,) rmem, one value per register (pre-transpose distribution). + dequant_scale: Optional[cutlass.Float32] = None, + dequant_scales: Optional[cute.Tensor] = None, ) -> cute.Tensor: # ── Input contract checks (compile-time): fp32, 1D, even-count, rmem ── # Wrapped in const_expr so the DSL evaluates them at trace time and the @@ -1770,8 +2027,9 @@ def alpha_swiglu_clamp( # gate_rmem / up_rmem are the RAW fc1 fp32 accumulator (pre-dequant). # Order follows the NVFP4 -> fp32 -> SwiGLU contract and MUST be: # - # 1. dequant: gate = alpha * gate_raw ; up = alpha * up_raw - # (alpha = expert-wise global scale on the acc; None => alpha == 1.) + # 1. dequant: gate = alpha * s_t * gate_raw ; up = alpha * s_t * up_raw + # (alpha = expert-wise global scale on the acc; s_t = per-token activation scale + # the input was divided by before quantization; None => 1 for either.) # 2. clamp the DEQUANTED (real) values, gpt-oss ``_apply_gate`` style: # gate = min(gate, +limit) (upper bound only) # up = clamp(up, -limit, +limit) (symmetric) @@ -1789,14 +2047,28 @@ def alpha_swiglu_clamp( out = cute.make_rmem_tensor((n,), cutlass.Float32) log2_e = 1.4426950408889634 - activation_alpha = 1.0 if self.swiglu_alpha is None else self.swiglu_alpha - neg_activation_alpha_log2e = -activation_alpha * log2_e - neg_log2e_pair = (cutlass.Float32(neg_activation_alpha_log2e), cutlass.Float32(neg_activation_alpha_log2e)) + # quick_geglu: out = up * gate * sigmoid(1.702 * gate). The 1.702 folds into the exp2 constant + # (one Python-side double product, rounded once to fp32), so the SwiGLU instruction sequence + # is reused unchanged. SiTU never reaches here with quick_geglu (_resolve_gate_activation). + sigmoid_scale = self._quick_geglu_sigmoid_scale if self.gate_activation == "quick_geglu" else 1.0 + if cutlass.const_expr(self.swiglu_alpha is not None): + sigmoid_scale = self.swiglu_alpha + neg_log2e_pair = (cutlass.Float32(-sigmoid_scale * log2_e), cutlass.Float32(-sigmoid_scale * log2_e)) one_pair = (cutlass.Float32(1.0), cutlass.Float32(1.0)) if cutlass.const_expr(self.swiglu_beta is not None): beta_pair = (cutlass.Float32(self.swiglu_beta), cutlass.Float32(self.swiglu_beta)) if cutlass.const_expr(self.gate_up_clamp is not None): limit = cutlass.Float32(self.gate_up_clamp) + if cutlass.const_expr(dequant_scale is not None and dequant_scales is not None): + raise ValueError("alpha_swiglu_clamp: pass at most one of dequant_scale and dequant_scales.") + if cutlass.const_expr(alpha_val is not None and (dequant_scale is not None or dequant_scales is not None)): + raise ValueError( + "alpha_swiglu_clamp: a dequant multiplier already carries alpha; pass alpha_val=None with it." + ) + if cutlass.const_expr(dequant_scales is not None and cute.size(dequant_scales) != n): + raise ValueError( + f"alpha_swiglu_clamp: dequant_scales must hold {n} values, got {cute.size(dequant_scales)}" + ) for i in cutlass.range_constexpr(0, n, 2): g0 = gate_rmem[i] @@ -1804,11 +2076,18 @@ def alpha_swiglu_clamp( u0 = up_rmem[i] u1 = up_rmem[i + 1] - # 1) dequant raw acc to real values (skip entirely when alpha is None). - if cutlass.const_expr(alpha_val is not None): - alpha_pair = (alpha_val, alpha_val) - g0, g1 = cute.arch.mul_packed_f32x2((g0, g1), alpha_pair) - u0, u1 = cute.arch.mul_packed_f32x2((u0, u1), alpha_pair) + # 1) dequant raw acc to real values (skip entirely when alpha and scale are None). + if cutlass.const_expr(dequant_scales is not None): + scale_pair = (dequant_scales[i], dequant_scales[i + 1]) + elif cutlass.const_expr(dequant_scale is not None): + scale_pair = (dequant_scale, dequant_scale) + elif cutlass.const_expr(alpha_val is not None): + scale_pair = (alpha_val, alpha_val) + else: + scale_pair = None + if cutlass.const_expr(scale_pair is not None): + g0, g1 = cute.arch.mul_packed_f32x2((g0, g1), scale_pair) + u0, u1 = cute.arch.mul_packed_f32x2((u0, u1), scale_pair) # 2) clamp the real values (skip when no clamp configured). if cutlass.const_expr(self.gate_up_clamp is not None): diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py index fc47fe0e3fa5..bcead8bcb700 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py @@ -14,7 +14,7 @@ from ....schedulers.fc12_mapping import BlockPhase, SwapAbFc12WorkTileInfo, peek_ready_bit -TensorRole = Literal["a", "b", "sfa", "sfb", "c", "sfc", "topk"] +TensorRole = Literal["a", "b", "sfa", "sfb", "c", "sfc", "topk", "fc1_act_scale"] @cute.jit @@ -155,7 +155,8 @@ def get_gmem_tensor( result = cute.domain_offset((data_token_offset, 0, 0), tensor) return (_rewrite_tensor_shape(result, (shape[0], shape[1], singleton)), None) - if cutlass.const_expr(tensor_name == "topk"): + if cutlass.const_expr(tensor_name in ("topk", "fc1_act_scale")): + # Both follow the data-token pool (not the independently padded SF pool). return (cute.domain_offset((data_token_offset,), tensor), None) raise ValueError(f"Unknown tensor_name: {tensor_name!r}.") diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_kernel.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_kernel.py index c60f1e814697..c35a024f48a7 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_kernel.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_kernel.py @@ -53,6 +53,11 @@ def problem_desc_require(cls) -> dict[str, object]: "gate_up_clamp": Optional[float], "situ_beta": OptionalRequirement(Optional[float]), "situ_linear_beta": OptionalRequirement(Optional[float]), + # "swiglu" (default) or "quick_geglu"; validated by the epilogue. + "gate_activation": OptionalRequirement(str), + # Whether __call__ receives fc1_activation_per_token_scale; reserves the epilogue's + # SMEM scratch for it. Absent means False. + "fc1_activation_per_token_scale": OptionalRequirement(bool), } @classmethod @@ -90,6 +95,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.gate_up_clamp = problem_desc["gate_up_clamp"] self.situ_beta = problem_desc.get("situ_beta") self.situ_linear_beta = problem_desc.get("situ_linear_beta") + self.gate_activation = problem_desc.get("gate_activation") or "swiglu" self.mma_tiler_mnk = impl_desc["mma_tiler_mnk"] self.cluster_shape_mn = impl_desc["cluster_shape_mn"] @@ -281,6 +287,10 @@ def __call__( fc1_alpha: Optional[cute.Tensor] = None, fc2_alpha: Optional[cute.Tensor] = None, fc1_norm_const: Optional[cute.Tensor] = None, + # (tokens,) fp32, same grouped layout as topk_scores. The activation was quantized from + # x / scale; the FC1 epilogue multiplies it back before the gated activation. Requires + # problem_desc["fc1_activation_per_token_scale"] = True (the epilogue rejects it otherwise). + fc1_activation_per_token_scale: Optional[cute.Tensor] = None, ) -> None: """Prepare and launch the lean fused FC1+FC2 kernel. @@ -289,6 +299,13 @@ def __call__( """ if cutlass.const_expr((expert_token_sizes is None) == (expert_token_prefix_sum is None)): raise ValueError("Exactly one of expert_token_sizes and expert_token_prefix_sum must be provided.") + if cutlass.const_expr( + fc1_activation_per_token_scale is not None + and fc1_activation_per_token_scale.element_type is not cutlass.Float32 + ): + raise TypeError( + f"fc1_activation_per_token_scale must be Float32, got {fc1_activation_per_token_scale.element_type}." + ) # The quant kind fixes every operand format, so the tensors are where a caller mismatch can # still show up. See BlockScaledSwapAbMegaMoeKernel.__call__. for operand_name, operand, expected_dtype in ( @@ -486,6 +503,7 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: fc2_alpha, fc1_norm_const, topk_scores, + fc1_activation_per_token_scale, ).launch( grid=grid, block=[self.threads_per_cta, 1, 1], @@ -526,6 +544,7 @@ def _kernel( fc2_alpha: Optional[cute.Tensor], fc1_norm_const: Optional[cute.Tensor], topk_scores: Optional[cute.Tensor], + fc1_activation_per_token_scale: Optional[cute.Tensor], ): """Compose Scheduler, Mainloop, and Epilogue.""" self._mainloop.materialize_codegen_members() @@ -599,7 +618,11 @@ def _kernel( fc1_ready_counter_pointer=None, ) optional_epilogue_args = GatedActEpilogueArgs( - fc1_alpha=fc1_alpha, fc2_alpha=fc2_alpha, fc1_norm_const=fc1_norm_const, topk_scores=topk_scores + fc1_alpha=fc1_alpha, + fc2_alpha=fc2_alpha, + fc1_norm_const=fc1_norm_const, + topk_scores=topk_scores, + fc1_activation_per_token_scale=fc1_activation_per_token_scale, ) pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py index 7dad5c86e02b..2e0497d3aa72 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py @@ -77,11 +77,17 @@ def problem_desc_require(cls) -> dict[str, object]: "swiglu_beta": OptionalRequirement(Optional[float]), "situ_beta": OptionalRequirement(Optional[float]), "situ_linear_beta": OptionalRequirement(Optional[float]), + # "swiglu" (default) or "quick_geglu"; validated by the epilogue, keyed in name(). + "gate_activation": OptionalRequirement(str), "world_size": int, "topk": int, "topk_index_dtype": type, "max_tokens_per_rank": int, "apply_topk_at_fc1": bool, + # The FC1 activation was quantized from x / s_t with a per-token fp32 s_t; the kernel + # multiplies s_t back before the gated activation. Optional so existing descriptors + # stay valid; it keys the compiled kernel (workspace region + epilogue path). + "fc1_activation_per_token_scale": OptionalRequirement(bool), } @classmethod @@ -95,6 +101,7 @@ def impl_desc_require(cls) -> dict[str, type]: "sf_padding_block": int, "work_id_mode": str, "fc2_use_bulk": bool, + "epi_async_release": OptionalRequirement(bool), "epi_flag_batches": tuple, "launch_cluster_count": int, "fallback_cluster_shape_mn": OptionalRequirement(Optional[tuple]), @@ -133,11 +140,13 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.swiglu_beta = problem_desc.get("swiglu_beta") self.situ_beta = problem_desc.get("situ_beta") self.situ_linear_beta = problem_desc.get("situ_linear_beta") + self.gate_activation = problem_desc.get("gate_activation") or "swiglu" self.world_size = problem_desc["world_size"] self.topk = problem_desc["topk"] self.topk_index_dtype = problem_desc["topk_index_dtype"] self.max_tokens_per_rank = problem_desc["max_tokens_per_rank"] self.apply_topk_at_fc1 = problem_desc["apply_topk_at_fc1"] + self.fc1_activation_per_token_scale = bool(problem_desc.get("fc1_activation_per_token_scale", False)) self.mma_tiler_mnk = impl_desc["mma_tiler_mnk"] self.cluster_shape_mn = impl_desc["cluster_shape_mn"] @@ -246,6 +255,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: "communication_enabled": True, "token_back_push_data": self.token_comm.token_back_push_data, "fc1_epi_flag_batch": self.epi_flag_batches[0], + "epi_async_release": impl_desc.get("epi_async_release", True), "fc2_epi_flag_batch": self.epi_flag_batches[1], } ) @@ -453,7 +463,7 @@ def dtype_name(dtype: type) -> str: f"sched{self.schedule_mode}_hint{self.hint}_pad{self.token_padding_block}x{self.sf_padding_block}_" f"work{self.work_id_mode}_" f"fc2store{'tma' if self.epilogue.fc2_use_tma else 'ublk' if self.epilogue.fc2_use_ublk else 'stg'}_" - f"fc2tmastages{self.epilogue.fc2_tma_stages}_" + f"fc2tmastages{self.epilogue.fc2_tma_stages}_asyncrel{int(self.epilogue.epi_async_release)}_" f"epiflag{epi_flags}_clusters{self.launch_cluster_count}_ctas{self.mixed_cga_config.total_cta_cnt}_" f"tokeninflag{self.token_in_flag_batch}_tokenback{self.token_back_mode}_" f"combine{self.combine_format.name}_clamp{self.gate_up_clamp}_" @@ -463,6 +473,9 @@ def dtype_name(dtype: type) -> str: f"{'apply_topk_fc1' if self.apply_topk_at_fc1 else 'apply_topk_fc2'}_" f"{'inkernel_redg' if self.reduce_topk_in_kernel else 'separate_reduce'}" "_dlb_v1" + # Off / SwiGLU keep the historical name so existing compiled-kernel caches stay valid. + + ("_fc1actscale" if self.fc1_activation_per_token_scale else "") + + (f"_act{self.gate_activation}" if self.gate_activation != "swiglu" else "") ) def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): @@ -550,6 +563,11 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): fc2_alpha=fake_tensor(cutlass.Float32, (experts,), (0,), set(), 4), fc1_norm_const=fake_tensor(cutlass.Float32, (experts,), (0,), set(), 4), ) + if self.fc1_activation_per_token_scale: + # Local (tokens,) fp32 read by the router on the source rank; only in the ABI when enabled. + fake_arguments.update( + fc1_activation_per_token_scale=fake_tensor(cutlass.Float32, (tokens,), (0,), {0}, 4), + ) compiled = cute.compile[cute.EnableTVMFFI(True)](self, **fake_arguments) if out_path is None: @@ -613,12 +631,40 @@ def __call__( fc1_norm_const: Optional[cute.Tensor] = None, # (local_experts,) hot_expert_weight_ready_flags: Optional[cute.Pointer] = None, # uint64[EP] hot_expert_weight_ready_generation: Optional[cutlass.Uint64] = None, + # (tokens,) fp32, this rank's own tokens, plain local memory (the router reads it here and + # pushes it to the pool). Required iff the descriptor enables fc1_activation_per_token_scale. + # Last on purpose: the AOT kwargs wrapper right-aligns its None defaults, so appending keeps + # every existing positional caller valid. + fc1_activation_per_token_scale: Optional[cute.Tensor] = None, ) -> None: """Launch router, fused MegaMoE compute, and optional TopK reduce.""" def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: return cute.make_tensor(tensor.iterator, cute.make_layout(shape, stride=tensor.stride)) + if cutlass.const_expr(self.fc1_activation_per_token_scale != (fc1_activation_per_token_scale is not None)): + raise ValueError( + "fc1_activation_per_token_scale must be passed iff the descriptor enables it " + f"(enabled={self.fc1_activation_per_token_scale})." + ) + if cutlass.const_expr( + fc1_activation_per_token_scale is not None + and fc1_activation_per_token_scale.element_type is not cutlass.Float32 + ): + raise TypeError( + f"fc1_activation_per_token_scale must be Float32, got {fc1_activation_per_token_scale.element_type}." + ) + if cutlass.const_expr(fc1_activation_per_token_scale is not None): + # The router gathers scale[token] for every routed token, so the scale needs one entry + # per topk_indices row. Only checkable when both extents are static (AOT leaves them + # dynamic); a short tensor would otherwise read past its end silently. + scale_rows = fc1_activation_per_token_scale.shape[0] + token_rows = topk_indices.shape[0] + if cutlass.const_expr(isinstance(scale_rows, int) and isinstance(token_rows, int) and scale_rows < token_rows): + raise ValueError( + f"fc1_activation_per_token_scale has {scale_rows} entries but topk_indices has {token_rows} rows." + ) + helper_ready_bundle = ( hot_expert_weight_ready_flags, hot_expert_weight_ready_generation, @@ -674,6 +720,7 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: peer_rank_ptr_mapper_host, self._device_workspace, stream, + fc1_activation_per_token_scale, ) peer_rank_ptr_mapper = peer_rank_ptr_mapper_host.make_device_object() @@ -980,6 +1027,7 @@ def _kernel( token_src_metadata = self.token_comm.token_src_metadata_tensor(self._device_workspace) fc2_done_counter = self.token_comm.fc2_done_counter_tensor(self._device_workspace) pool_topk_scores = self.token_comm.fc1_topk_scores_tensor(self._device_workspace) + pool_activation_scales = self.token_comm.fc1_activation_per_token_scales_tensor(self._device_workspace) # Bind the active cluster before deriving the AB pipeline arrival count. self._mainloop.assign_device_members( @@ -1028,7 +1076,11 @@ def _kernel( fc1_ready_counter_pointer=self.token_comm.fc1_ready_counter_pointer(self._device_workspace), ) optional_epilogue_args = GatedActEpilogueArgs( - fc1_alpha=fc1_alpha, fc2_alpha=fc2_alpha, fc1_norm_const=fc1_norm_const, topk_scores=pool_topk_scores + fc1_alpha=fc1_alpha, + fc2_alpha=fc2_alpha, + fc1_norm_const=fc1_norm_const, + topk_scores=pool_topk_scores, + fc1_activation_per_token_scale=pool_activation_scales, ) pipeline_init_arrive_mixed_cga( diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py index 7fbb2be22b76..911ced925fb0 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py @@ -71,6 +71,11 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.swiglu_alpha = float(self.swiglu_alpha) self.swiglu_beta = float(self.swiglu_beta) self.situ_beta, self.situ_linear_beta = self._resolve_situ_betas(problem_desc) + # quick_geglu is Blackwell-only so far; refuse rather than silently computing SwiGLU. + if problem_desc.get("gate_activation") not in (None, "swiglu"): + raise NotImplementedError( + f"Rubin local epilogue does not support gate_activation={problem_desc.get('gate_activation')!r} yet." + ) self.mma_tiler_mnk = impl_desc["mma_tiler_mnk"] self.cluster_shape_mn = impl_desc["cluster_shape_mn"] self.use_2cta_instrs = impl_desc["use_2cta_instrs"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py index f0868f743ba4..2d503d43c52f 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py @@ -960,6 +960,12 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.swiglu_alpha = float(self.swiglu_alpha) self.swiglu_beta = float(self.swiglu_beta) self.situ_beta, self.situ_linear_beta = self._resolve_situ_betas(problem_desc) + # The Blackwell epilogue accepts gate_activation="quick_geglu"; this one has not been + # ported yet. Refuse rather than silently computing SwiGLU for a quick_geglu descriptor. + if problem_desc.get("gate_activation") not in (None, "swiglu"): + raise NotImplementedError( + f"Rubin epilogue does not support gate_activation={problem_desc.get('gate_activation')!r} yet." + ) self.mma_tiler_mnk = impl_desc["mma_tiler_mnk"] self.cluster_shape_mn = impl_desc["cluster_shape_mn"] self.use_2cta_instrs = impl_desc["use_2cta_instrs"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py index c5567f3df5a0..9d99afc3344a 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py @@ -19,10 +19,9 @@ and the FC2 epilogue warps send results to their owning rank themselves through the bulk reduce-add. Twelve warps cover every role. -The kernel owns the workspace tail reset. Because it outlives the communication -kernel, it is the only participant that can safely zero the counters, which it -does between two NVLink barriers so that no rank observes a half-cleared -workspace or exits before the clearing is visible. +The kernel owns the workspace tail reset. All warp roles finish before a grid +barrier allows the epilogue to clear the counters. A final NVLink barrier makes +the reset visible to every rank before the next launch reuses the workspace. """ import hashlib @@ -1210,31 +1209,6 @@ def _kernel( tmem_allocator.relinquish_alloc_permit() tmem_allocator.free(tmem_allocator.retrieve_ptr(self.acc_dtype), self._mainloop.num_tmem_alloc_cols) - epilogue_thread_count = 32 * len(self.epilogue_warp_ids) - - iket.range_push("tail.grid_sync_before_reset") - self.tail_barrier.sync( - epilogue_thread_count, Int32(self.promised_launchable_sm_count), linear_cta_idx, thread_idx - ) - iket.range_pop() - - iket.range_push("tail.reset_workspace") - total_reset_threads = self.promised_launchable_sm_count * epilogue_thread_count - global_reset_thread = linear_cta_idx * Int32(epilogue_thread_count) + thread_idx - self._device_workspace.reset_tail_space("shared", global_reset_thread, total_reset_threads) - self._device_workspace.reset_tail_space("local", global_reset_thread, total_reset_threads) - iket.range_pop() - - iket.range_push("tail.nvlink_publish") - self.tail_barrier.arrive_and_wait( - epilogue_thread_count, - Int32(self.promised_launchable_sm_count), - linear_cta_idx, - thread_idx, - prologue_grid_sync=True, - epilogue_grid_sync=False, - ) - iket.range_pop() iket.range_pop() if (warp_idx >= Int32(self.gather_b_warp_ids[0])) & (warp_idx <= Int32(self.gather_b_warp_ids[-1])): @@ -1265,6 +1239,36 @@ def _kernel( ) iket.range_pop() + # CTAs with no work can finish the epilogue before gather warps pass + # their readiness waits. Join every role before clearing those flags. + cute.arch.sync_threads() + if warp_idx < len(self.epilogue_warp_ids): + epilogue_thread_count = 32 * len(self.epilogue_warp_ids) + + iket.range_push("tail.grid_sync_before_reset") + self.tail_barrier.sync( + epilogue_thread_count, Int32(self.promised_launchable_sm_count), linear_cta_idx, thread_idx + ) + iket.range_pop() + + iket.range_push("tail.reset_workspace") + total_reset_threads = self.promised_launchable_sm_count * epilogue_thread_count + global_reset_thread = linear_cta_idx * Int32(epilogue_thread_count) + thread_idx + self._device_workspace.reset_tail_space("shared", global_reset_thread, total_reset_threads) + self._device_workspace.reset_tail_space("local", global_reset_thread, total_reset_threads) + iket.range_pop() + + iket.range_push("tail.nvlink_publish") + self.tail_barrier.arrive_and_wait( + epilogue_thread_count, + Int32(self.promised_launchable_sm_count), + linear_cta_idx, + thread_idx, + prologue_grid_sync=True, + epilogue_grid_sync=False, + ) + iket.range_pop() + self.tail_barrier.remove_device_members() self._device_workspace.remove_device_members() From 76bfad0b6631a81266a1480090b2e78927a3303c Mon Sep 17 00:00:00 2001 From: jintaop Date: Mon, 19 Oct 2026 07:00:37 -0700 Subject: [PATCH 05/59] [None][chore] Drop the internal-GitLab MegaMoE vendor lock entries Remove the cutedsl-megamoe and megamoe-scheduler entries from 3rdparty/vendor_sources.lock.yaml, together with 3rdparty/vendor_patches/cutedsl-megamoe.patch. Both entries pointed at internal GitLab repositories, which the public repository should not reference. The lock is now identical to main. The sources under tensorrt_llm/_torch/cute_dsl_kernels/cutedsl_megamoe and megamoe_scheduler are unchanged and are maintained in-tree, as cutedsl_megamoe already is on main. They come from cutedsl_megamoe main 7832a330 plus the TensorRT-LLM edits the patch carried (SwiGLU alpha/beta, bandit B105 annotations, the #19960 fc12_scheduler change and one comment edit), and from the HALO-Q scheduler at 4dc86160. The ruff exclusion comment no longer says the sources are pinned. Signed-off-by: jintaop --- 3rdparty/vendor_patches/cutedsl-megamoe.patch | 919 ------------------ 3rdparty/vendor_sources.lock.yaml | 33 - pyproject.toml | 2 +- 3 files changed, 1 insertion(+), 953 deletions(-) delete mode 100644 3rdparty/vendor_patches/cutedsl-megamoe.patch diff --git a/3rdparty/vendor_patches/cutedsl-megamoe.patch b/3rdparty/vendor_patches/cutedsl-megamoe.patch deleted file mode 100644 index 9c594f14ec2c..000000000000 --- a/3rdparty/vendor_patches/cutedsl-megamoe.patch +++ /dev/null @@ -1,919 +0,0 @@ -diff --git a/communication/nvlink_domain/token_comm.py b/communication/nvlink_domain/token_comm.py -index a64a5e7e377b77373863a766db4a1ed70c9b9252..374aa01b055bb187dd588c64ec0482a54c9d1d32 100644 ---- a/communication/nvlink_domain/token_comm.py -+++ b/communication/nvlink_domain/token_comm.py -@@ -123,7 +123,7 @@ class _MetadataPushRouter(KernelComponent): - sorted_metadata_region = "nvlink.token_comm.sorted_metadata" - sorted_scores_region = "nvlink.token_comm.sorted_scores" - pool_expert_base_region = "nvlink.token_comm.pool_expert_base" -- token_src_metadata_region = "nvlink.token_comm.token_src_metadata" -+ token_src_metadata_region = "nvlink.token_comm.token_src_metadata" # nosec B105 - fc1_topk_scores_region = "nvlink.token_comm.fc1_topk_scores" - # Pool-ordered fp32 per-token FC1 activation scale, pushed next to the metadata like the scores. - fc1_activation_scales_region = "nvlink.token_comm.fc1_activation_per_token_scales" -@@ -1178,16 +1178,16 @@ class TokenCommNonDeterministic(KernelComponent): - fc2_activation_sf_region = "nvlink.token_comm.fc2_activation_sf" - pre_reduced_activation_region = "nvlink.token_comm.pre_reduced_activation" - pre_reduced_activation_sf_region = "nvlink.token_comm.pre_reduced_activation_sf" -- token_back_schedule_region = "nvlink.token_comm.token_back_schedule" -+ token_back_schedule_region = "nvlink.token_comm.token_back_schedule" # nosec B105 - -- token_in_mbarrier_region = "nvlink.token_comm.main_smem.token_in_mbarriers" -- token_back_mbarrier_region = "nvlink.token_comm.main_smem.token_back_mbarriers" -+ token_in_mbarrier_region = "nvlink.token_comm.main_smem.token_in_mbarriers" # nosec B105 -+ token_back_mbarrier_region = "nvlink.token_comm.main_smem.token_back_mbarriers" # nosec B105 - expert_sizes_smem_region = "nvlink.token_comm.main_smem.expert_sizes" -- token_in_activation_smem_region = "nvlink.token_comm.main_smem.token_in_activation" -- token_in_sf_smem_region = "nvlink.token_comm.main_smem.token_in_sf" -- token_in_sf_group_region = "nvlink.token_comm.main_smem.token_in_sf_group" -- token_back_activation_smem_region = "nvlink.token_comm.main_smem.token_back_activation" -- token_back_sf_smem_region = "nvlink.token_comm.main_smem.token_back_sf" -+ token_in_activation_smem_region = "nvlink.token_comm.main_smem.token_in_activation" # nosec B105 -+ token_in_sf_smem_region = "nvlink.token_comm.main_smem.token_in_sf" # nosec B105 -+ token_in_sf_group_region = "nvlink.token_comm.main_smem.token_in_sf_group" # nosec B105 -+ token_back_activation_smem_region = "nvlink.token_comm.main_smem.token_back_activation" # nosec B105 -+ token_back_sf_smem_region = "nvlink.token_comm.main_smem.token_back_sf" # nosec B105 - - @classmethod - def problem_desc_require(cls) -> dict[str, type]: -@@ -1277,7 +1277,7 @@ class TokenCommNonDeterministic(KernelComponent): - # FC12 occupies warps 0..7. Legacy primitive callers retain their - # modulo-four numbering by omitting the explicit CTA ordinal. - self.dispatch_warp_start = 8 -- self.combine_warp_start = 12 if self.token_back_mode == "standalone_warps" else 8 -+ self.combine_warp_start = 12 if self.token_back_mode == "standalone_warps" else 8 # nosec B105 - self.dispatch_warp_count = self.transfer_warp_count - self.combine_warp_count = self.transfer_warp_count - self.dispatch_cta_count = self.promised_launchable_sm_count -@@ -1356,9 +1356,9 @@ class TokenCommNonDeterministic(KernelComponent): - ) - if self.token_back_ready_granularity not in ("expert", "token_tile"): - raise ValueError(f"Unsupported token_back_ready_granularity {self.token_back_ready_granularity!r}.") -- if self.token_back_ready_granularity == "token_tile" and ( -+ if self.token_back_ready_granularity == "token_tile" and ( # nosec B105 - self.token_back_mode not in ("standalone_warps", "reuse_dispatch_warps") -- or self.token_back_schedule_mode != "atomic_counter" -+ or self.token_back_schedule_mode != "atomic_counter" # nosec B105 - ): - raise ValueError("token_tile return-ready requires standalone_warps or reuse_dispatch_warps with atomic_counter.") - if not 1 <= self.token_in_flag_batch <= 32: -@@ -1391,7 +1391,7 @@ class TokenCommNonDeterministic(KernelComponent): - raise ValueError("token_in_stage_split must cut a token row into whole 16-byte chunks.") - - if self.token_back_enabled and self.combine_stage_depth > 1: -- if self.token_back_schedule_mode != "static": -+ if self.token_back_schedule_mode != "static": # nosec B105 - raise ValueError("Multi-stage token-back requires static scheduling.") - if self.combine_format.is_quantized: - raise ValueError("Multi-stage token-back supports BF16 combine only.") -@@ -1438,7 +1438,7 @@ class TokenCommNonDeterministic(KernelComponent): - - @property - def token_back_push_data(self) -> bool: -- return self.token_back_mode != "epi_warps" -+ return self.token_back_mode != "epi_warps" # nosec B105 - - @property - def token_back_push_sf(self) -> bool: -@@ -1663,7 +1663,7 @@ class TokenCommNonDeterministic(KernelComponent): - if self.token_back_enabled: - fc2_done_slots = ( - self.max_fc1_ready_slot_count -- if self.token_back_ready_granularity == "token_tile" -+ if self.token_back_ready_granularity == "token_tile" # nosec B105 - else self.experts_per_rank - ) - workspace.register( -@@ -1687,7 +1687,7 @@ class TokenCommNonDeterministic(KernelComponent): - buffer_space="local", - byte_alignment=128, - ) -- if self.token_back_enabled and self.token_back_schedule_mode == "atomic_counter": -+ if self.token_back_enabled and self.token_back_schedule_mode == "atomic_counter": # nosec B105 - workspace.register( - self.token_back_schedule_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" - ) -@@ -1732,7 +1732,7 @@ class TokenCommNonDeterministic(KernelComponent): - self._register_token_in_staging(transfer_overlay.add_lifetime("token_in")) - if not self.token_back_enabled: - return -- if self.token_back_mode == "standalone_warps": -+ if self.token_back_mode == "standalone_warps": # nosec B105 - back_overlay = workspace.create_overlay("nvlink.token_comm.main_smem.standalone_token_back") - else: - back_overlay = transfer_overlay -@@ -1788,7 +1788,7 @@ class TokenCommNonDeterministic(KernelComponent): - combine_rings = self.combine_warp_count - combine_depth = self.combine_stage_depth - sf_padded = self.activation_sf_hidden_padded -- if self.token_back_mode == "standalone_warps": -+ if self.token_back_mode == "standalone_warps": # nosec B105 - available_bytes_per_warp = self.token_back_chunk_bytes - else: - sf_bytes = sf_padded * int(self.activation_sf_dtype.width) // 8 if self.activation_is_quantized else 0 -@@ -1796,13 +1796,13 @@ class TokenCommNonDeterministic(KernelComponent): - - if self.token_back_push_data: - bytes_per_output_token = self.hidden_size * int(self.combine_format.act_dtype.width) // 8 -- if self.token_back_mode == "standalone_warps": -+ if self.token_back_mode == "standalone_warps": # nosec B105 - chunk_bytes = self.token_back_chunk_bytes - elif available_bytes_per_warp < bytes_per_output_token: - chunk_bytes = self.bytes_per_token - else: - chunk_bytes = bytes_per_output_token -- if self.token_back_mode != "standalone_warps" and bytes_per_output_token % chunk_bytes != 0: -+ if self.token_back_mode != "standalone_warps" and bytes_per_output_token % chunk_bytes != 0: # nosec B105 - raise ValueError("Token-back data chunk bytes must divide one row.") - chunk_elements = chunk_bytes * 8 // int(self.combine_format.act_dtype.width) - shape = ( -@@ -2303,7 +2303,7 @@ class TokenCommNonDeterministic(KernelComponent): - ) - token_in_size_barrier.arrive_and_wait() - iket.range_pop() -- if cutlass.const_expr(not self.cta_specialized and self.token_back_mode == "standalone_warps"): -+ if cutlass.const_expr(not self.cta_specialized and self.token_back_mode == "standalone_warps"): # nosec B105 - sizes_ready_barrier = pipeline.NamedBarrier( - barrier_id=self.standalone_size_barrier_id, - num_threads=2 * self.transfer_thread_count, -@@ -2495,7 +2495,7 @@ class TokenCommNonDeterministic(KernelComponent): - - flag_tracker.fire() - cute.arch.sync_warp() -- if cutlass.const_expr(self.token_back_enabled and self.token_back_mode != "standalone_warps"): -+ if cutlass.const_expr(self.token_back_enabled and self.token_back_mode != "standalone_warps"): # nosec B105 - iket.range_push("token_in.transfer_barrier") - transfer_lifetime_barrier = pipeline.NamedBarrier( - barrier_id=self.transfer_lifetime_barrier_id, -@@ -2621,7 +2621,7 @@ class TokenCommNonDeterministic(KernelComponent): - cute.arch.mbarrier_init(worker_mbarrier + Int32(init_stage), 1) - cute.arch.sync_warp() - -- if cutlass.const_expr(not self.cta_specialized and self.token_back_mode == "standalone_warps"): -+ if cutlass.const_expr(not self.cta_specialized and self.token_back_mode == "standalone_warps"): # nosec B105 - sizes_ready_barrier = pipeline.NamedBarrier( - barrier_id=self.standalone_size_barrier_id, - num_threads=2 * self.transfer_thread_count, -@@ -2819,7 +2819,7 @@ class TokenCommNonDeterministic(KernelComponent): - cute.arch.mbarrier_init(worker_mbarrier, 1) - cute.arch.sync_warp() - -- if cutlass.const_expr(not self.cta_specialized and self.token_back_mode == "standalone_warps"): -+ if cutlass.const_expr(not self.cta_specialized and self.token_back_mode == "standalone_warps"): # nosec B105 - sizes_ready_barrier = pipeline.NamedBarrier( - barrier_id=self.standalone_size_barrier_id, - num_threads=2 * self.transfer_thread_count, -@@ -2835,7 +2835,7 @@ class TokenCommNonDeterministic(KernelComponent): - activation_chunk_count = ceil_div(output_token_bytes, activation_chunk_bytes) - data_window_unit = activation_chunk_bytes * 2 - reuse_data_pacing_enabled = ( -- self.token_back_mode == "reuse_dispatch_warps" and data_window_unit > self.minimum_pacing_window_cycles -+ self.token_back_mode == "reuse_dispatch_warps" and data_window_unit > self.minimum_pacing_window_cycles # nosec B105 - ) - # Preserve the empirical low:initial:high ratio of 1:2.5:5. - data_average_window = Int32(data_window_unit) -@@ -2847,19 +2847,19 @@ class TokenCommNonDeterministic(KernelComponent): - sf_chunk_count = ceil_div(output_sf_bytes, sf_chunk_bytes) - sf_window_unit = ceil_div(sf_chunk_bytes * 2, 3) - reuse_sf_pacing_enabled = ( -- self.token_back_mode == "reuse_dispatch_warps" and sf_window_unit > self.minimum_pacing_window_cycles -+ self.token_back_mode == "reuse_dispatch_warps" and sf_window_unit > self.minimum_pacing_window_cycles # nosec B105 - ) - sf_average_window = Int32(sf_window_unit) - sf_low_window = sf_window_unit * 2 // 5 - sf_high_window = sf_window_unit * 2 - - next_dense_token = global_worker_idx - global_worker_count -- if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): -+ if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 - next_dense_token = Int32(0) - next_dense_token = self.next_token(next_dense_token) - expert_valid_begin = Int32(0) - transfer_phase = Int32(0) -- if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): -+ if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): # nosec B105 - expert_ready_slot_begin = Int32(0) - last_ready_slot = Int32(-1) - -@@ -2867,7 +2867,7 @@ class TokenCommNonDeterministic(KernelComponent): - while local_expert < Int32(self.experts_per_rank): - expert_token_count = owned_sizes[local_expert] - expert_valid_end = expert_valid_begin + expert_token_count -- if cutlass.const_expr(self.token_back_ready_granularity == "expert"): -+ if cutlass.const_expr(self.token_back_ready_granularity == "expert"): # nosec B105 - if next_dense_token < expert_valid_end: - token_tile_count = (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) // Int32( - self.tokens_per_fc1_ready_slot -@@ -2880,7 +2880,7 @@ class TokenCommNonDeterministic(KernelComponent): - - while next_dense_token < expert_valid_end: - token_in_expert = next_dense_token - expert_valid_begin -- if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): -+ if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): # nosec B105 - # Same rank-local expert/tile prefix as FC1 ready; separate counter storage. - ready_slot = expert_ready_slot_begin + token_in_expert // Int32(self.tokens_per_fc1_ready_slot) - if ready_slot != last_ready_slot: -@@ -2925,7 +2925,7 @@ class TokenCommNonDeterministic(KernelComponent): - current_chunk_bytes = Int32(chunk_bytes_this_round) - current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) - stateless_data_pacing_enabled = ( -- self.token_back_mode != "reuse_dispatch_warps" -+ self.token_back_mode != "reuse_dispatch_warps" # nosec B105 - and current_window_unit > self.minimum_pacing_window_cycles - ) - round_start_clock = Int64(0) -@@ -3001,7 +3001,7 @@ class TokenCommNonDeterministic(KernelComponent): - current_chunk_bytes = Int32(chunk_bytes_this_round) - current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) - stateless_sf_pacing_enabled = ( -- self.token_back_mode != "reuse_dispatch_warps" -+ self.token_back_mode != "reuse_dispatch_warps" # nosec B105 - and current_window_unit > self.minimum_pacing_window_cycles - ) - round_start_clock = Int64(0) -@@ -3054,7 +3054,7 @@ class TokenCommNonDeterministic(KernelComponent): - cute.arch.sync_warp() - next_dense_token = self.next_token(next_dense_token) - -- if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): -+ if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): # nosec B105 - # Advance even when this worker skipped the expert (empty experts contribute zero). - expert_ready_slot_begin = expert_ready_slot_begin + ( - expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1) -@@ -3070,9 +3070,9 @@ class TokenCommNonDeterministic(KernelComponent): - # communication CTAs' combine warps, not every persistent CTA's transfer warps. - global_worker_count = self.combine_cta_count * self.combine_warp_count - schedule_counter = None -- if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): -+ if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 - schedule_counter = self._device_workspace.ptr(self.token_back_schedule_region) -- if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): -+ if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 - claimed_token = Int32(0) - if lane_idx == Int32(0): - claimed_token = cute.arch.atomic_add(schedule_counter, Int32(1), sem="relaxed", scope="gpu") -diff --git a/communication/nvlink_domain/token_comm_cta_mainloop.py b/communication/nvlink_domain/token_comm_cta_mainloop.py -index e9cf01a60538d05649f80d79c0cf8cac129566f7..ecc713fcbf9df7cba4ca999b45eee3732167db33 100644 ---- a/communication/nvlink_domain/token_comm_cta_mainloop.py -+++ b/communication/nvlink_domain/token_comm_cta_mainloop.py -@@ -29,7 +29,7 @@ class TokenCommCtaMainloop: - self.dispatch_warp_count = token_comm.dispatch_warp_count - self.combine_warp_start = token_comm.combine_warp_start - self.combine_warp_count = token_comm.combine_warp_count -- if token_comm.token_back_mode != "standalone_warps": -+ if token_comm.token_back_mode != "standalone_warps": # nosec B105 - raise ValueError("Dedicated communication CTAs require standalone token-back warps.") - if not token_back_on_compute_ctas: - dispatch_end = self.dispatch_warp_start + self.dispatch_warp_count -diff --git a/communication/nvlink_domain/token_comm_deterministic.py b/communication/nvlink_domain/token_comm_deterministic.py -index 5f0f38f79326d591d9f1103bcc182cb95729c28f..6838edb87603bd90c56c76a47114582d5baf869a 100644 ---- a/communication/nvlink_domain/token_comm_deterministic.py -+++ b/communication/nvlink_domain/token_comm_deterministic.py -@@ -138,7 +138,7 @@ class _MetadataPushRouter(KernelComponent): - sorted_metadata_region = "nvlink.token_comm.sorted_metadata" - sorted_scores_region = "nvlink.token_comm.sorted_scores" - pool_expert_base_region = "nvlink.token_comm.pool_expert_base" -- token_src_metadata_region = "nvlink.token_comm.token_src_metadata" -+ token_src_metadata_region = "nvlink.token_comm.token_src_metadata" # nosec B105 - fc1_topk_scores_region = "nvlink.token_comm.fc1_topk_scores" - source_expert_base_region = "nvlink.token_comm.source_expert_base" - push_destination_base_region = "nvlink.token_comm.push_destination_base" -@@ -1375,15 +1375,15 @@ class TokenCommDeterministic(KernelComponent): - fc2_activation_sf_region = "nvlink.token_comm.fc2_activation_sf" - pre_reduced_activation_region = "nvlink.token_comm.pre_reduced_activation" - pre_reduced_activation_sf_region = "nvlink.token_comm.pre_reduced_activation_sf" -- token_back_schedule_region = "nvlink.token_comm.token_back_schedule" -+ token_back_schedule_region = "nvlink.token_comm.token_back_schedule" # nosec B105 - -- token_in_mbarrier_region = "nvlink.token_comm.main_smem.token_in_mbarriers" -- token_back_mbarrier_region = "nvlink.token_comm.main_smem.token_back_mbarriers" -+ token_in_mbarrier_region = "nvlink.token_comm.main_smem.token_in_mbarriers" # nosec B105 -+ token_back_mbarrier_region = "nvlink.token_comm.main_smem.token_back_mbarriers" # nosec B105 - expert_sizes_smem_region = "nvlink.token_comm.main_smem.expert_sizes" -- token_in_activation_smem_region = "nvlink.token_comm.main_smem.token_in_activation" -- token_in_sf_smem_region = "nvlink.token_comm.main_smem.token_in_sf" -- token_back_activation_smem_region = "nvlink.token_comm.main_smem.token_back_activation" -- token_back_sf_smem_region = "nvlink.token_comm.main_smem.token_back_sf" -+ token_in_activation_smem_region = "nvlink.token_comm.main_smem.token_in_activation" # nosec B105 -+ token_in_sf_smem_region = "nvlink.token_comm.main_smem.token_in_sf" # nosec B105 -+ token_back_activation_smem_region = "nvlink.token_comm.main_smem.token_back_activation" # nosec B105 -+ token_back_sf_smem_region = "nvlink.token_comm.main_smem.token_back_sf" # nosec B105 - - @classmethod - def problem_desc_require(cls) -> dict[str, type]: -@@ -1573,7 +1573,7 @@ class TokenCommDeterministic(KernelComponent): - - @property - def token_back_push_data(self) -> bool: -- return self.token_back_mode != "epi_warps" -+ return self.token_back_mode != "epi_warps" # nosec B105 - - @property - def token_back_push_sf(self) -> bool: -@@ -1776,7 +1776,7 @@ class TokenCommDeterministic(KernelComponent): - buffer_space="local", - byte_alignment=128, - ) -- if self.token_back_enabled and self.token_back_schedule_mode == "atomic_counter": -+ if self.token_back_enabled and self.token_back_schedule_mode == "atomic_counter": # nosec B105 - workspace.register( - self.token_back_schedule_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" - ) -@@ -1825,14 +1825,14 @@ class TokenCommDeterministic(KernelComponent): - if not self.token_back_enabled: - return - -- if self.token_back_mode == "standalone_warps": -+ if self.token_back_mode == "standalone_warps": # nosec B105 - token_back_lifetime = workspace.create_overlay( - "nvlink.token_comm.main_smem.standalone_token_back" - ).add_lifetime("token_back") - else: - token_back_lifetime = transfer_overlay.add_lifetime("token_back") - -- if self.token_back_mode == "standalone_warps": -+ if self.token_back_mode == "standalone_warps": # nosec B105 - available_bytes_per_warp = self.standalone_chunk_bytes - else: - activation_bytes = self.bytes_per_token -@@ -1845,13 +1845,13 @@ class TokenCommDeterministic(KernelComponent): - - if self.token_back_push_data: - bytes_per_output_token = self.hidden_size * int(self.combine_format.act_dtype.width) // 8 -- if self.token_back_mode == "standalone_warps": -+ if self.token_back_mode == "standalone_warps": # nosec B105 - chunk_bytes = self.standalone_chunk_bytes - elif available_bytes_per_warp < bytes_per_output_token: - chunk_bytes = self.bytes_per_token - else: - chunk_bytes = bytes_per_output_token -- if self.token_back_mode != "standalone_warps" and bytes_per_output_token % chunk_bytes != 0: -+ if self.token_back_mode != "standalone_warps" and bytes_per_output_token % chunk_bytes != 0: # nosec B105 - raise ValueError("Token-back data chunk bytes must divide one row.") - chunk_elements = chunk_bytes * 8 // int(self.combine_format.act_dtype.width) - token_back_lifetime.register_tensor( -@@ -2175,7 +2175,7 @@ class TokenCommDeterministic(KernelComponent): - ) - token_in_size_barrier.arrive_and_wait() - iket.range_pop() -- if cutlass.const_expr(self.token_back_mode == "standalone_warps"): -+ if cutlass.const_expr(self.token_back_mode == "standalone_warps"): # nosec B105 - sizes_ready_barrier = pipeline.NamedBarrier( - barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count - ) -@@ -2364,7 +2364,7 @@ class TokenCommDeterministic(KernelComponent): - flag_tracker.fire() - cute.arch.sync_warp() - iket.range_pop() -- if cutlass.const_expr(self.token_back_enabled and self.token_back_mode != "standalone_warps"): -+ if cutlass.const_expr(self.token_back_enabled and self.token_back_mode != "standalone_warps"): # nosec B105 - iket.range_push("token_in.transfer_barrier") - transfer_lifetime_barrier = pipeline.NamedBarrier( - barrier_id=self.transfer_lifetime_barrier_id, num_threads=self.transfer_thread_count -@@ -2432,7 +2432,7 @@ class TokenCommDeterministic(KernelComponent): - cute.arch.sync_warp() - - owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) -- if cutlass.const_expr(self.token_back_mode == "standalone_warps"): -+ if cutlass.const_expr(self.token_back_mode == "standalone_warps"): # nosec B105 - sizes_ready_barrier = pipeline.NamedBarrier( - barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count - ) -@@ -2447,7 +2447,7 @@ class TokenCommDeterministic(KernelComponent): - activation_chunk_count = ceil_div(output_token_bytes, activation_chunk_bytes) - data_window_unit = ceil_div(activation_chunk_bytes * 2, 3) - reuse_data_pacing_enabled = ( -- self.token_back_mode == "reuse_dispatch_warps" and data_window_unit > self.minimum_pacing_window_cycles -+ self.token_back_mode == "reuse_dispatch_warps" and data_window_unit > self.minimum_pacing_window_cycles # nosec B105 - ) - # Preserve the empirical low:initial:high ratio of 1:2.5:5. - data_average_window = Int32(data_window_unit) -@@ -2460,14 +2460,14 @@ class TokenCommDeterministic(KernelComponent): - sf_chunk_count = ceil_div(output_sf_bytes, sf_chunk_bytes) - sf_window_unit = ceil_div(sf_chunk_bytes * 2, 3) - reuse_sf_pacing_enabled = ( -- self.token_back_mode == "reuse_dispatch_warps" and sf_window_unit > self.minimum_pacing_window_cycles -+ self.token_back_mode == "reuse_dispatch_warps" and sf_window_unit > self.minimum_pacing_window_cycles # nosec B105 - ) - sf_average_window = Int32(sf_window_unit) - sf_low_window = sf_window_unit * 2 // 5 - sf_high_window = sf_window_unit * 2 - - next_dense_token = global_worker_idx - global_worker_count -- if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): -+ if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 - next_dense_token = Int32(0) - next_dense_token = self.next_token(next_dense_token) - expert_valid_begin = Int32(0) -@@ -2523,7 +2523,7 @@ class TokenCommDeterministic(KernelComponent): - current_chunk_bytes = Int32(chunk_bytes_this_round) - current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) - stateless_data_pacing_enabled = ( -- self.token_back_mode != "reuse_dispatch_warps" -+ self.token_back_mode != "reuse_dispatch_warps" # nosec B105 - and current_window_unit > self.minimum_pacing_window_cycles - ) - round_start_clock = Int64(0) -@@ -2590,7 +2590,7 @@ class TokenCommDeterministic(KernelComponent): - current_chunk_bytes = Int32(chunk_bytes_this_round) - current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) - stateless_sf_pacing_enabled = ( -- self.token_back_mode != "reuse_dispatch_warps" -+ self.token_back_mode != "reuse_dispatch_warps" # nosec B105 - and current_window_unit > self.minimum_pacing_window_cycles - ) - round_start_clock = Int64(0) -@@ -2642,9 +2642,9 @@ class TokenCommDeterministic(KernelComponent): - def next_token(self, current_token: Int32) -> Int32: - global_worker_count = self.promised_launchable_sm_count * self.transfer_warp_count - schedule_counter = None -- if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): -+ if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 - schedule_counter = self._device_workspace.ptr(self.token_back_schedule_region) -- if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): -+ if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 - claimed_token = Int32(0) - if self._lane_idx == Int32(0): - claimed_token = cute.arch.atomic_add(schedule_counter, Int32(1), sem="relaxed", scope="gpu") -diff --git a/helpers/topk_reduce.py b/helpers/topk_reduce.py -index 6e4d0fe7d363b8554e5b2789240e555ccc2499ea..ce4cb6465a0fb430bcb2a93acf2d8a158ff2989b 100644 ---- a/helpers/topk_reduce.py -+++ b/helpers/topk_reduce.py -@@ -38,7 +38,7 @@ from .dsl_helpers import mark_alignment - # fp4 (e2m1) -> fp32 register decode. - # - # Blackwell has no e2m1->f32 upconvert: the framework's ``term.load().to(f32)`` --# lowers to an ALU subnormal-normalization path (~60% DRAM SOL). Both helpers -+# lowers to an ALU subnormal-normalization path. Both helpers - # below force a table-driven decode instead; ``e2m1_reg`` (N e2m1 codes, N % 8 - # == 0) is read as packed b32 words and N fp32 values are written into - # ``fp32_reg`` in code order. The 16 e2m1 values are exact in fp32, so the two -diff --git a/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py b/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py -index a30dc98d4dd1bbae7600a6d3894fed8e7eec7c18..a363fcdd831a043315607dc277d0345706f80e6a 100644 ---- a/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py -+++ b/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py -@@ -968,6 +968,8 @@ class SwapABGatedActEpilogue(KernelComponent): - "intermediate_gateup_size": StaticOrRuntimeIntegerType, - "combine_format": CombineFormat, - "gate_up_clamp": Optional[float], -+ "swiglu_alpha": OptionalRequirement(Optional[float]), -+ "swiglu_beta": OptionalRequirement(Optional[float]), - # SiTU (Kimi K3) selects a different gated-activation core; optional so - # existing SwiGLU descriptors stay valid unchanged. See _resolve_situ_betas. - "situ_beta": OptionalRequirement(Optional[float]), -@@ -1018,6 +1020,8 @@ class SwapABGatedActEpilogue(KernelComponent): - raise ValueError(f"SiTU beta parameters must be positive, got {beta} and {linear_beta}.") - if problem_desc["gate_up_clamp"] is not None: - raise ValueError("SiTU does not support gate_up_clamp; the two activation variants are exclusive.") -+ if problem_desc.get("swiglu_alpha") is not None: -+ raise ValueError("SiTU does not support SwiGLUBias constants.") - return float(beta), float(linear_beta) - - @classmethod -@@ -1037,6 +1041,8 @@ class SwapABGatedActEpilogue(KernelComponent): - ) - if gate_activation != "swiglu" and problem_desc.get("situ_beta") is not None: - raise ValueError(f"gate_activation={gate_activation!r} cannot be combined with SiTU betas.") -+ if gate_activation != "swiglu" and problem_desc.get("swiglu_alpha") is not None: -+ raise ValueError(f"gate_activation={gate_activation!r} cannot be combined with SwiGLUBias constants.") - return gate_activation - - def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: -@@ -1048,6 +1054,13 @@ class SwapABGatedActEpilogue(KernelComponent): - self.intermediate_gateup_size = problem_desc["intermediate_gateup_size"] - self.combine_format = problem_desc["combine_format"] - self.gate_up_clamp = problem_desc["gate_up_clamp"] -+ self.swiglu_alpha = problem_desc.get("swiglu_alpha") -+ self.swiglu_beta = problem_desc.get("swiglu_beta") -+ if (self.swiglu_alpha is None) != (self.swiglu_beta is None): -+ raise ValueError("swiglu_alpha and swiglu_beta must be set together.") -+ if self.swiglu_alpha is not None: -+ self.swiglu_alpha = float(self.swiglu_alpha) -+ self.swiglu_beta = float(self.swiglu_beta) - self.situ_beta, self.situ_linear_beta = self._resolve_situ_betas(problem_desc) - self.gate_activation = self._resolve_gate_activation(problem_desc) - # The scratch costs mainloop SMEM budget, so only a descriptor that enables the scale -@@ -2021,7 +2034,8 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): - # gate = min(gate, +limit) (upper bound only) - # up = clamp(up, -limit, +limit) (symmetric) - # 3. gated activation, either SwiGLU or SiTU (mutually exclusive; SiTU has no clamp): -- # SwiGLU: out = up * gate * sigmoid(gate) -+ # SwiGLU: out = (up + swiglu_beta) * gate * sigmoid(swiglu_alpha * gate) -+ # Standard SwiGLU uses swiglu_alpha=1 and swiglu_beta=0. - # SiTU: out = beta*tanh(gate/beta)*sigmoid(gate) * linear_beta*tanh(up/linear_beta) - # sigmoid(x) = rcp(1 + exp2(-x * log2e)) - # -@@ -2037,8 +2051,12 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): - # (one Python-side double product, rounded once to fp32), so the SwiGLU instruction sequence - # is reused unchanged. SiTU never reaches here with quick_geglu (_resolve_gate_activation). - sigmoid_scale = self._quick_geglu_sigmoid_scale if self.gate_activation == "quick_geglu" else 1.0 -+ if cutlass.const_expr(self.swiglu_alpha is not None): -+ sigmoid_scale = self.swiglu_alpha - neg_log2e_pair = (cutlass.Float32(-sigmoid_scale * log2_e), cutlass.Float32(-sigmoid_scale * log2_e)) - one_pair = (cutlass.Float32(1.0), cutlass.Float32(1.0)) -+ if cutlass.const_expr(self.swiglu_beta is not None): -+ beta_pair = (cutlass.Float32(self.swiglu_beta), cutlass.Float32(self.swiglu_beta)) - if cutlass.const_expr(self.gate_up_clamp is not None): - limit = cutlass.Float32(self.gate_up_clamp) - if cutlass.const_expr(dequant_scale is not None and dequant_scales is not None): -@@ -2131,7 +2149,10 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): - sigmoid_pair = _sigmoid(g0, g1) - - if cutlass.const_expr(self.situ_beta is None): -- # SwiGLU: out = up * gate * sigmoid(gate) -+ # SwiGLU: out = (up + swiglu_beta) * gate * sigmoid(swiglu_alpha * gate) -+ # The SwiGLUBias offset follows the up clamp. -+ if cutlass.const_expr(self.swiglu_beta is not None): -+ u0, u1 = cute.arch.add_packed_f32x2((u0, u1), beta_pair) - ug = cute.arch.mul_packed_f32x2((u0, u1), (g0, g1)) - out_pair = cute.arch.mul_packed_f32x2(ug, sigmoid_pair) - else: -diff --git a/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py b/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py -index ae93a57df562296bc79d08c4594c8c3cd9ba09aa..2e0497d3aa720edd03717fd0117fb5fd8d735499 100644 ---- a/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py -+++ b/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py -@@ -73,6 +73,8 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): - "b_major_mode": OperandMajorMode, - "combine_format": CombineFormat, - "gate_up_clamp": Optional[float], -+ "swiglu_alpha": OptionalRequirement(Optional[float]), -+ "swiglu_beta": OptionalRequirement(Optional[float]), - "situ_beta": OptionalRequirement(Optional[float]), - "situ_linear_beta": OptionalRequirement(Optional[float]), - # "swiglu" (default) or "quick_geglu"; validated by the epilogue, keyed in name(). -@@ -134,6 +136,8 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): - self.b_major_mode = problem_desc["b_major_mode"] - self.combine_format = problem_desc["combine_format"] - self.gate_up_clamp = problem_desc["gate_up_clamp"] -+ self.swiglu_alpha = problem_desc.get("swiglu_alpha") -+ self.swiglu_beta = problem_desc.get("swiglu_beta") - self.situ_beta = problem_desc.get("situ_beta") - self.situ_linear_beta = problem_desc.get("situ_linear_beta") - self.gate_activation = problem_desc.get("gate_activation") or "swiglu" -@@ -177,8 +181,8 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): - # static M is required only when an S>0 helper boundary must be - # compiled into the scheduler gate. - self.home_expert_count = self.local_expert_count -- self.threads_per_cta = 16 * 32 if self.token_back_mode == "standalone_warps" else 12 * 32 -- self.other_warp_register_count = 64 if self.token_back_mode == "standalone_warps" else 72 -+ self.threads_per_cta = 16 * 32 if self.token_back_mode == "standalone_warps" else 12 * 32 # nosec B105 -+ self.other_warp_register_count = 64 if self.token_back_mode == "standalone_warps" else 72 # nosec B105 - - self.mixed_cga_config = NonClcMixedCgaConfig( - preferred_cluster_shape=self.cluster_shape_mn, -@@ -463,6 +467,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): - f"epiflag{epi_flags}_clusters{self.launch_cluster_count}_ctas{self.mixed_cga_config.total_cta_cnt}_" - f"tokeninflag{self.token_in_flag_batch}_tokenback{self.token_back_mode}_" - f"combine{self.combine_format.name}_clamp{self.gate_up_clamp}_" -+ f"swiglua{self.swiglu_alpha}_swiglub{self.swiglu_beta}_" - # The SiTU betas are codegen-time constants, so they must key the compiled kernel. - f"situ{self.situ_beta}x{self.situ_linear_beta}_" - f"{'apply_topk_fc1' if self.apply_topk_at_fc1 else 'apply_topk_fc2'}_" -@@ -1084,7 +1089,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): - pipeline_init_wait_mixed_cga( - self.cluster_shape_mn, self.resolved_fallback_cluster_shape_mn, is_fallback_cluster - ) -- finalize_barrier_thread_count = 12 * 32 if self.token_back_mode == "standalone_warps" else 8 * 32 -+ finalize_barrier_thread_count = 12 * 32 if self.token_back_mode == "standalone_warps" else 8 * 32 # nosec B105 - finalize_barrier = pipeline.NamedBarrier(barrier_id=13, num_threads=finalize_barrier_thread_count) - - if warp_idx == self.scheduler_warp_id: -@@ -1213,7 +1218,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): - self.token_comm.token_in(self._smem_workspace, smem_base) - iket.range_pop() - if cutlass.const_expr( -- self.token_comm.token_back_enabled and self.token_back_mode != "standalone_warps" -+ self.token_comm.token_back_enabled and self.token_back_mode != "standalone_warps" # nosec B105 - ): - iket.range_push("mega.token_back") - self.token_comm.token_back(self._smem_workspace, smem_base) -@@ -1222,7 +1227,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): - iket.range_push("mega.tail_reset") - self.token_comm.reset_tail() - iket.range_pop() -- elif cutlass.const_expr(self.token_back_mode == "standalone_warps"): -+ elif cutlass.const_expr(self.token_back_mode == "standalone_warps"): # nosec B105 - iket.range_push("mega.token_back_standalone") - self.token_comm.token_back(self._smem_workspace, smem_base) - iket.range_pop() -diff --git a/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py b/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py -index a97e10dd651fd728fd1e73dd3ec2791ee1b4ef9a..911ced925fb0d789d19983ec68e817cb36a0c329 100644 ---- a/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py -+++ b/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py -@@ -63,6 +63,13 @@ class SwapABGatedActEpilogue(_SharedSwapABGatedActEpilogue): - self.intermediate_gateup_size = problem_desc["intermediate_gateup_size"] - self.combine_format = problem_desc["combine_format"] - self.gate_up_clamp = problem_desc["gate_up_clamp"] -+ self.swiglu_alpha = problem_desc.get("swiglu_alpha") -+ self.swiglu_beta = problem_desc.get("swiglu_beta") -+ if (self.swiglu_alpha is None) != (self.swiglu_beta is None): -+ raise ValueError("swiglu_alpha and swiglu_beta must be set together.") -+ if self.swiglu_alpha is not None: -+ self.swiglu_alpha = float(self.swiglu_alpha) -+ self.swiglu_beta = float(self.swiglu_beta) - self.situ_beta, self.situ_linear_beta = self._resolve_situ_betas(problem_desc) - # quick_geglu is Blackwell-only so far; refuse rather than silently computing SwiGLU. - if problem_desc.get("gate_activation") not in (None, "swiglu"): -diff --git a/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_local_mega_moe_kernel.py b/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_local_mega_moe_kernel.py -index a486110e0e3eb4b673d5857f7d08484cc24893a5..3e86b1ad34e2e31ab8b3a394208e921156c4b552 100644 ---- a/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_local_mega_moe_kernel.py -+++ b/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_local_mega_moe_kernel.py -@@ -67,6 +67,8 @@ class BlockScaledSwapAbLocalMegaMoeKernel(KernelClass): - "a_major_mode": OperandMajorMode, - "b_major_mode": OperandMajorMode, - "gate_up_clamp": Optional[float], -+ "swiglu_alpha": OptionalRequirement(Optional[float]), -+ "swiglu_beta": OptionalRequirement(Optional[float]), - "situ_beta": OptionalRequirement(Optional[float]), - "situ_linear_beta": OptionalRequirement(Optional[float]), - "topk": int, -@@ -116,6 +118,8 @@ class BlockScaledSwapAbLocalMegaMoeKernel(KernelClass): - self.b_major_mode = problem_desc["b_major_mode"] - self.combine_format = CombineFormat.parse("bf16") - self.gate_up_clamp = problem_desc["gate_up_clamp"] -+ self.swiglu_alpha = problem_desc.get("swiglu_alpha") -+ self.swiglu_beta = problem_desc.get("swiglu_beta") - self.situ_beta = problem_desc.get("situ_beta") - self.situ_linear_beta = problem_desc.get("situ_linear_beta") - self.topk = problem_desc["topk"] -@@ -410,6 +414,7 @@ class BlockScaledSwapAbLocalMegaMoeKernel(KernelClass): - f"fc1epiflag{self.fc1_epi_flag_batch}_clusters{self.launch_cluster_count}_" - f"routingflag{self.routing_flag_batch}_" - f"combine{self.combine_format.name}_clamp{self.gate_up_clamp}_" -+ f"swiglua{self.swiglu_alpha}_swiglub{self.swiglu_beta}_" - # The SiTU betas are codegen-time constants, so they must key the compiled kernel. - f"situ{self.situ_beta}x{self.situ_linear_beta}_" - f"{'apply_topk_fc1' if self.apply_topk_at_fc1 else 'apply_topk_fc2'}_" -diff --git a/kernel_src/rubin/inference/local_mega/local_routing.py b/kernel_src/rubin/inference/local_mega/local_routing.py -index ce7877ec231429f6487668527285388c09bea0bd..fd3ceb087e89c02c041b0e794dc22fbc15a1e67b 100644 ---- a/kernel_src/rubin/inference/local_mega/local_routing.py -+++ b/kernel_src/rubin/inference/local_mega/local_routing.py -@@ -174,7 +174,7 @@ class _MetadataPushRouter(KernelComponent): - sorted_metadata_region = "local_mega.routing.sorted_metadata" - sorted_scores_region = "local_mega.routing.sorted_scores" - pool_expert_base_region = "local_mega.routing.pool_expert_base" -- token_src_metadata_region = "local_mega.routing.token_src_metadata" -+ token_src_metadata_region = "local_mega.routing.token_src_metadata" # nosec B105 - fc1_topk_scores_region = "local_mega.routing.fc1_topk_scores" - source_expert_base_region = "local_mega.routing.source_expert_base" - sorted_metadata_ready_region = "local_mega.routing.sorted_metadata_ready" -@@ -1021,9 +1021,9 @@ class LocalRouting(KernelComponent): - fc1_activation_sf_region = "local_mega.routing.fc1_activation_sf" - pre_reduced_activation_region = "local_mega.routing.pre_reduced_activation" - -- token_in_mbarrier_region = "local_mega.routing.main_smem.token_in_mbarriers" -+ token_in_mbarrier_region = "local_mega.routing.main_smem.token_in_mbarriers" # nosec B105 - expert_sizes_smem_region = "local_mega.routing.main_smem.expert_sizes" -- token_in_sf_smem_region = "local_mega.routing.main_smem.token_in_sf" -+ token_in_sf_smem_region = "local_mega.routing.main_smem.token_in_sf" # nosec B105 - - @classmethod - def problem_desc_require(cls) -> dict[str, type]: -diff --git a/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py b/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py -index 21f00a540720eff436d738b706819ef85502f023..2d503d43c52ff7ef9e04a00a6dc05c993b07662f 100644 ---- a/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py -+++ b/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py -@@ -896,6 +896,8 @@ class SwapABGatedActEpilogue(KernelComponent): - "intermediate_gateup_size": StaticOrRuntimeIntegerType, - "combine_format": CombineFormat, - "gate_up_clamp": Optional[float], -+ "swiglu_alpha": OptionalRequirement(Optional[float]), -+ "swiglu_beta": OptionalRequirement(Optional[float]), - # SiTU (Kimi K3) selects a different gated-activation core; optional so - # existing SwiGLU descriptors stay valid unchanged. See _resolve_situ_betas. - "situ_beta": OptionalRequirement(Optional[float]), -@@ -937,6 +939,8 @@ class SwapABGatedActEpilogue(KernelComponent): - raise ValueError(f"SiTU beta parameters must be positive, got {beta} and {linear_beta}.") - if problem_desc["gate_up_clamp"] is not None: - raise ValueError("SiTU does not support gate_up_clamp; the two activation variants are exclusive.") -+ if problem_desc.get("swiglu_alpha") is not None: -+ raise ValueError("SiTU does not support SwiGLUBias constants.") - return float(beta), float(linear_beta) - - def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: -@@ -948,6 +952,13 @@ class SwapABGatedActEpilogue(KernelComponent): - self.intermediate_gateup_size = problem_desc["intermediate_gateup_size"] - self.combine_format = problem_desc["combine_format"] - self.gate_up_clamp = problem_desc["gate_up_clamp"] -+ self.swiglu_alpha = problem_desc.get("swiglu_alpha") -+ self.swiglu_beta = problem_desc.get("swiglu_beta") -+ if (self.swiglu_alpha is None) != (self.swiglu_beta is None): -+ raise ValueError("swiglu_alpha and swiglu_beta must be set together.") -+ if self.swiglu_alpha is not None: -+ self.swiglu_alpha = float(self.swiglu_alpha) -+ self.swiglu_beta = float(self.swiglu_beta) - self.situ_beta, self.situ_linear_beta = self._resolve_situ_betas(problem_desc) - # The Blackwell epilogue accepts gate_activation="quick_geglu"; this one has not been - # ported yet. Refuse rather than silently computing SwiGLU for a quick_geglu descriptor. -@@ -1645,7 +1656,8 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): - # gate = min(gate, +limit) (upper bound only) - # up = clamp(up, -limit, +limit) (symmetric) - # 3. gated activation, either SwiGLU or SiTU (mutually exclusive; SiTU has no clamp): -- # SwiGLU: out = up * gate * sigmoid(gate) -+ # SwiGLU: out = (up + swiglu_beta) * gate * sigmoid(swiglu_alpha * gate) -+ # Standard SwiGLU uses swiglu_alpha=1 and swiglu_beta=0. - # SiTU: out = beta*tanh(gate/beta)*sigmoid(gate) * linear_beta*tanh(up/linear_beta) - # sigmoid(x) = rcp(1 + exp2(-x * log2e)) - # -@@ -1657,8 +1669,12 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): - out = cute.make_rmem_tensor((n,), cutlass.Float32) - log2_e = 1.4426950408889634 - -- neg_log2e_pair = (cutlass.Float32(-log2_e), cutlass.Float32(-log2_e)) -+ activation_alpha = 1.0 if self.swiglu_alpha is None else self.swiglu_alpha -+ neg_activation_alpha_log2e = -activation_alpha * log2_e -+ neg_log2e_pair = (cutlass.Float32(neg_activation_alpha_log2e), cutlass.Float32(neg_activation_alpha_log2e)) - one_pair = (cutlass.Float32(1.0), cutlass.Float32(1.0)) -+ if cutlass.const_expr(self.swiglu_beta is not None): -+ beta_pair = (cutlass.Float32(self.swiglu_beta), cutlass.Float32(self.swiglu_beta)) - if cutlass.const_expr(self.gate_up_clamp is not None): - limit = cutlass.Float32(self.gate_up_clamp) - -@@ -1734,7 +1750,10 @@ class SwapABFc1Epilogue(_ImmutableAfterInit): - sigmoid_pair = _sigmoid(g0, g1) - - if cutlass.const_expr(self.situ_beta is None): -- # SwiGLU: out = up * gate * sigmoid(gate) -+ # SwiGLU: out = (up + swiglu_beta) * gate * sigmoid(swiglu_alpha * gate) -+ # The SwiGLUBias offset follows the up clamp. -+ if cutlass.const_expr(self.swiglu_beta is not None): -+ u0, u1 = cute.arch.add_packed_f32x2((u0, u1), beta_pair) - ug = cute.arch.mul_packed_f32x2((u0, u1), (g0, g1)) - out_pair = cute.arch.mul_packed_f32x2(ug, sigmoid_pair) - else: -@@ -1991,7 +2010,7 @@ class SwapABFc2Epilogue(_ImmutableAfterInit): - publish: cutlass.Constexpr = self.token_back_enabled - flag_address = Int64(0) - if cutlass.const_expr(publish): -- if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): -+ if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): # nosec B105 - # The scheduler maps every fallback CTA into the preferred - # logical cluster. N=1 keeps this slot common to all hidden - # shards; padded hidden CTAs must also publish their completion. -diff --git a/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py b/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py -index f49de382daeef5a9751a19dd190008f02827c54c..3c5bb77f999cd0013dd68ec11668adefe816727d 100644 ---- a/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py -+++ b/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py -@@ -75,6 +75,8 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): - "b_major_mode": OperandMajorMode, - "combine_format": CombineFormat, - "gate_up_clamp": Optional[float], -+ "swiglu_alpha": OptionalRequirement(Optional[float]), -+ "swiglu_beta": OptionalRequirement(Optional[float]), - "situ_beta": OptionalRequirement(Optional[float]), - "situ_linear_beta": OptionalRequirement(Optional[float]), - "world_size": int, -@@ -136,6 +138,8 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): - self.b_major_mode = problem_desc["b_major_mode"] - self.combine_format = problem_desc["combine_format"] - self.gate_up_clamp = problem_desc["gate_up_clamp"] -+ self.swiglu_alpha = problem_desc.get("swiglu_alpha") -+ self.swiglu_beta = problem_desc.get("swiglu_beta") - self.situ_beta = problem_desc.get("situ_beta") - self.situ_linear_beta = problem_desc.get("situ_linear_beta") - self.world_size = problem_desc["world_size"] -@@ -202,8 +206,8 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): - # static M is required only when an S>0 helper boundary must be - # compiled into the scheduler gate. - self.home_expert_count = self.local_expert_count -- self.threads_per_cta = 16 * 32 if self.token_back_mode == "standalone_warps" else 12 * 32 -- self.other_warp_register_count = 64 if self.token_back_mode == "standalone_warps" else 72 -+ self.threads_per_cta = 16 * 32 if self.token_back_mode == "standalone_warps" else 12 * 32 # nosec B105 -+ self.other_warp_register_count = 64 if self.token_back_mode == "standalone_warps" else 72 # nosec B105 - - self.mixed_cga_config = NonClcMixedCgaConfig( - preferred_cluster_shape=self.cluster_shape_mn, -@@ -239,7 +243,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): - sf_vec_size=self.sf_vec_size, - ) - tmem_plan = make_tcgen05_tmem_plan(mma_instruction, self.architecture, self.mma_tiler_mnk) -- if self.token_back_ready_granularity == "token_tile": -+ if self.token_back_ready_granularity == "token_tile": # nosec B105 - tokens_per_fc1_ready_slot, fc2_done_signals_per_token_tile = derive_fc2_ready_geometry( - hidden_size=self.hidden_size, - mma_tiler_mnk=self.mma_tiler_mnk, -@@ -301,7 +305,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): - ) - self.scheduler = scheduler_type(fc12_problem_desc, resolved_impl_desc) - self.epilogue = SwapABGatedActEpilogue(fc12_problem_desc, resolved_impl_desc) -- if self.token_back_ready_granularity == "token_tile": -+ if self.token_back_ready_granularity == "token_tile": # nosec B105 - self._validate_token_tile_ready_contract() - if self.epilogue.fc1_output_dtype is not self.b_dtype: - raise ValueError("Epilogue FC1 output dtype must match the Mainloop B dtype.") -@@ -366,7 +370,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): - or any(isinstance(w, bool) or not isinstance(w, int) or w <= 0 for w in warps) - ): - raise ValueError(f"A split needs comm_warps as two positive ints, got {warps}.") -- if self.token_back_mode != "standalone_warps": -+ if self.token_back_mode != "standalone_warps": # nosec B105 - raise ValueError("A CTA split requires token_back_mode='standalone_warps'.") - dispatch_warps, back_warps = warps - # Check the warps that actually coexist in each concrete CTA branch. -@@ -438,7 +442,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): - def _validate_geometry(self) -> None: - if self.token_back_ready_granularity not in ("expert", "token_tile"): - raise ValueError(f"Unsupported token_back_ready_granularity {self.token_back_ready_granularity!r}.") -- if self.token_back_ready_granularity == "token_tile": -+ if self.token_back_ready_granularity == "token_tile": # nosec B105 - if not ( - self.schedule_mode in ("grouped", "phase_interleave") - and self.work_id_mode == "atomic_counter" -@@ -678,12 +682,13 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): - f"epiflag{epi_flags}_clusters{self.launch_cluster_count}_ctas{self.mixed_cga_config.total_cta_cnt}_" - f"tokeninflag{self.token_in_flag_batch}_tokenback{self.token_back_mode}_" - f"combine{self.combine_format.name}_clamp{self.gate_up_clamp}_" -+ f"swiglua{self.swiglu_alpha}_swiglub{self.swiglu_beta}_" - # The SiTU betas are codegen-time constants, so they must key the compiled kernel. - f"situ{self.situ_beta}x{self.situ_linear_beta}_" - f"{'apply_topk_fc1' if self.apply_topk_at_fc1 else 'apply_topk_fc2'}_" - f"{'inkernel_redg' if self.reduce_topk_in_kernel else 'separate_reduce'}" - "_dlb_v1" -- + ("_fc2readytile_v1" if self.token_back_ready_granularity == "token_tile" else "") -+ + ("_fc2readytile_v1" if self.token_back_ready_granularity == "token_tile" else "") # nosec B105 - ) - - def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): -@@ -1275,7 +1280,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): - pipeline_init_wait_mixed_cga( - self.cluster_shape_mn, self.resolved_fallback_cluster_shape_mn, is_fallback_cluster - ) -- finalize_barrier_thread_count = 12 * 32 if self.token_back_mode == "standalone_warps" else 8 * 32 -+ finalize_barrier_thread_count = 12 * 32 if self.token_back_mode == "standalone_warps" else 8 * 32 # nosec B105 - finalize_barrier = pipeline.NamedBarrier(barrier_id=13, num_threads=finalize_barrier_thread_count) - - if warp_idx == self.scheduler_warp_id: -@@ -1405,7 +1410,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): - self.token_comm.token_in(self._smem_workspace, smem_base) - iket.range_pop() - if cutlass.const_expr( -- self.token_comm.token_back_enabled and self.token_back_mode != "standalone_warps" -+ self.token_comm.token_back_enabled and self.token_back_mode != "standalone_warps" # nosec B105 - ): - iket.range_push("mega.token_back") - self.token_comm.token_back(self._smem_workspace, smem_base) -@@ -1414,7 +1419,7 @@ class BlockScaledSwapAbMegaMoeKernel(KernelClass): - iket.range_push("mega.tail_reset") - self.token_comm.reset_tail() - iket.range_pop() -- elif cutlass.const_expr(self.token_back_mode == "standalone_warps"): -+ elif cutlass.const_expr(self.token_back_mode == "standalone_warps"): # nosec B105 - iket.range_push("mega.token_back_standalone") - self.token_comm.token_back(self._smem_workspace, smem_base) - iket.range_pop() -diff --git a/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py b/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py -index 043f2f86d7e9f0017ca28144f3d56b85eac4aa7b..9d99afc3344a55a00544fddbe3979dd0b0bf3032 100644 ---- a/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py -+++ b/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py -@@ -108,6 +108,8 @@ class BlockScaledSwapAbGenphaseMoeKernel(KernelClass): - "quant_kind": str, - "combine_format": CombineFormat, - "gate_up_clamp": Optional[float], -+ "swiglu_alpha": OptionalRequirement(Optional[float]), -+ "swiglu_beta": OptionalRequirement(Optional[float]), - "world_size": int, - "topk": int, - "topk_index_dtype": type, -@@ -168,6 +170,8 @@ class BlockScaledSwapAbGenphaseMoeKernel(KernelClass): - self.acc_dtype = tcgen05_block_scaled_acc_dtype - self.combine_format = problem_desc["combine_format"] - self.gate_up_clamp = problem_desc["gate_up_clamp"] -+ self.swiglu_alpha = problem_desc.get("swiglu_alpha") -+ self.swiglu_beta = problem_desc.get("swiglu_beta") - self.world_size = problem_desc["world_size"] - self.topk = problem_desc["topk"] - self.topk_index_dtype = problem_desc["topk_index_dtype"] -@@ -545,6 +549,7 @@ class BlockScaledSwapAbGenphaseMoeKernel(KernelClass): - f"epiflag{epi_flags}_clusters{self.launch_cluster_count}_" - f"pusher{self.pusher_cta_count}_refine{self.refine_participant_groups}x{self.refine_output_stages}_" - f"combine{self.combine_format.name}_clamp{self.gate_up_clamp}_" -+ f"swiglua{self.swiglu_alpha}_swiglub{self.swiglu_beta}_" - f"{'apply_topk_fc1' if self.apply_topk_at_fc1 else 'apply_topk_fc2'}_" - f"{'inkernel_reduce' if self.reduce_topk_in_kernel else 'separate_reduce'}_" - f"abstages{self._mainloop.num_weight_ab_stages}x{self._mainloop.num_token_ab_stages}" -diff --git a/kernel_src/schedulers/fc12_scheduler.py b/kernel_src/schedulers/fc12_scheduler.py -index b014b25f72a48e6be38922236faa0a84c9672a1a..9fd349cd80481271f08825c66ae28c53313ba93f 100644 ---- a/kernel_src/schedulers/fc12_scheduler.py -+++ b/kernel_src/schedulers/fc12_scheduler.py -@@ -906,8 +906,6 @@ class PhaseInterleavedFc12Scheduler(SchedulerBase): - minimum_claim_count, task_mapping_state - ) - spin_wait(fc1_counter_pointer, lambda value: value >= claim_target) -- if stream_ends_before_target: -- fc1_exhausted = Boolean(True) - return fc1_exhausted - - @cute.jit -diff --git a/kernel_src/schedulers/non_clc_mixed_cga.py b/kernel_src/schedulers/non_clc_mixed_cga.py -index 841157953fa0eac1ab379feecf0aed589055bc14..ba879cc6ac4d11bacaa4341f6e2f08765a4b1633 100644 ---- a/kernel_src/schedulers/non_clc_mixed_cga.py -+++ b/kernel_src/schedulers/non_clc_mixed_cga.py -@@ -134,7 +134,7 @@ class NonClcMixedCgaSchedulerWorker: - cluster_broadcast_region = "non_clc_mixed_cga.cluster_broadcast" - work_id_counter_region = "non_clc_mixed_cga.work_id_counters" - fallback_registration_counter_region = "non_clc_mixed_cga.fallback_registration_counter" -- fallback_group_token_region = "non_clc_mixed_cga.fallback_group_tokens" -+ fallback_group_token_region = "non_clc_mixed_cga.fallback_group_tokens" # nosec B105 - - def __init__( - self, diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 23d66e54e12b..18640486388d 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -1,26 +1,5 @@ schema_version: 1 vendors: - cutedsl-megamoe: - url: https://gitlab-master.nvidia.com/bangyus/cutedsl_megamoe.git - branch: main - commit: 7832a3300e24d791f94b09b3f6b37b5429e2fdf5 - source: sources - destination: tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe - include: - - '*.py' - - communication/**/*.py - - helpers/[!mt]*.py - - helpers/megamoe_aot.py - - helpers/topk_reduce.py - - kernel_src/*.py - - kernel_src/blackwell/*.py - - kernel_src/rubin/*.py - - kernel_src/*/inference/**/*.py - - kernel_src/schedulers/[!w]*.py - - kernel_src/schedulers/work_id_claim.py - patch: 3rdparty/vendor_patches/cutedsl-megamoe.patch - patch_digest: sha256:f24f6c5d296410bf3393ee5aafbb967982eee373a6e5f3289aed1b801b9b9d96 - digest: sha256-tree-v1:6e83535fe953947cc2b3c0e26a22cd8ded0dba1ff7845965e1f6b7ae36b2aac3 flashinfer-kda-bt16: url: https://github.com/flashinfer-ai/flashinfer.git branch: main @@ -44,18 +23,6 @@ vendors: patch: 3rdparty/vendor_patches/flashinfer-prims-ts.patch patch_digest: sha256:fc5c23ec3d65c562025a1845f211fdac66bfcd3da156dbb8385f4c4407d20d07 digest: sha256-tree-v1:f2f444f7e6d4ad688fbf4e1e9eed11ff7c78d835cfd0af5ef9d7660780b059e0 - megamoe-scheduler: - url: https://gitlab-master.nvidia.com/jintaop/cutedsl_eplb_scheduler_copy.git - branch: trtllm/pr-19646-single-plane-spec - commit: 4dc8616031b18dd948bdf11b49afe0d86be15e22 - source: megamoe_scheduler - destination: tensorrt_llm/_torch/moe/kernels/megamoe_scheduler - include: - - __init__.py - - cuda_scheduler/*.py - - geometry.py - - sami/*.py - digest: sha256-tree-v1:00210684fbffaa623fc00839bcacce000650db9090cf7a1a07bdd0ef3a9afc33 openengine-proto: url: https://github.com/ai-dynamo/openengine.git tag: v0.1.0 diff --git a/pyproject.toml b/pyproject.toml index ea7b9f5d70cb..6a73230989b2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,7 +49,7 @@ exclude = [ "tensorrt_llm/_torch/moe/kernels/rubin/manual_mma_128dp.py", "tensorrt_llm/_torch/moe/kernels/rubin/rubin_contiguous_grouped_blockscaled_gemm_fused_fc12.py", "tensorrt_llm/grpc/openengine/_generated/**", - # Keep pinned MegaMoE sources byte-identical to upstream. + # Keep vendored MegaMoE sources in upstream formatting. "tensorrt_llm/_torch/cute_dsl_kernels/cutedsl_megamoe/**", "tensorrt_llm/_torch/cute_dsl_kernels/megamoe_scheduler/**", ] From 95d0458da6f00f40a7c4c575151e8ac44bc58b75 Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 02:16:01 -0700 Subject: [PATCH 06/59] [None][chore] Follow the cute_dsl_kernels split on main Rebased onto main after #19275 split tensorrt_llm/_torch/cute_dsl_kernels/ into domain kernel trees. - Move the MegaMoE kernel sources to tensorrt_llm/_torch/moe/kernels/ (cutedsl_megamoe, megamoe_scheduler, megamoe_kernels.py, megamoe_shared_fc12.py); their contents are unchanged, except the autotuner import in megamoe_shared_fc12.py, which is one level deeper now. - Update the imports, the ruff and pre-commit excludes, and the contracts test to the new module paths. - Move test_quantization_sm_budget.py to tests/unittest/_torch/kernels/parallel/ and test_megamoe_shared_fc12_autotune.py to tests/unittest/_torch/moe/. - The MoE developer guide now describes cutedsl_megamoe as an in-tree copy of upstream instead of a lock-file vendor. Signed-off-by: jintaop --- .pre-commit-config.yaml | 2 +- pyproject.toml | 4 +- .../ops/moe/__init__.py | 27 +-------- .../kernels/rubin/dsv4_qb_fusion/__init__.py | 24 +------- .../_torch/modules/megamoe_shared_mlp.py | 2 +- .../moe/fused_moe/MOE_DEVELOPER_GUIDE.md | 6 +- .../mega_moe/rebalance_live_arena.py | 10 ++-- .../mega_moe/rebalance_slot_scheduler.py | 8 +-- .../_torch/moe/fused_moe/quantization.py | 3 +- .../cutedsl_megamoe/communication/__init__.py | 59 ++++++------------- .../communication/nvlink_domain/__init__.py | 41 +++++++------ .../cutedsl_megamoe/helpers/__init__.py | 59 +++++++++++++------ .../cutedsl_megamoe/kernel_src/__init__.py | 36 ++++++----- .../kernel_src/blackwell/__init__.py | 18 ++---- .../blackwell/inference/__init__.py | 23 +++++++- .../blackwell/inference/mega/__init__.py | 30 +++++----- .../kernel_src/rubin/__init__.py | 7 +-- .../kernel_src/rubin/inference/__init__.py | 25 ++++---- .../rubin/inference/local_mega/__init__.py | 31 +--------- .../rubin/inference/mega/__init__.py | 27 ++++++++- .../kernel_src/schedulers/__init__.py | 33 ++++++++++- .../_torch/moe/kernels/megamoe_shared_fc12.py | 2 +- .../integration/test_lists/test-db/l0_cpu.yml | 2 +- .../parallel/test_quantization_sm_budget.py | 0 .../_torch/moe/test_dynamic_eplb_contracts.py | 2 +- .../test_megamoe_shared_fc12_autotune.py | 2 +- 26 files changed, 240 insertions(+), 243 deletions(-) rename tests/unittest/_torch/{moe => }/kernels/parallel/test_quantization_sm_budget.py (100%) rename tests/unittest/_torch/{cute_dsl_kernels => moe}/test_megamoe_shared_fc12_autotune.py (99%) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index b2e83942a02d..100a078ef993 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1463,7 +1463,7 @@ legacy-files: &legacy_files | # list; the hook's own `files:` pattern only gates *when* the hook triggers. # Global exclude: vendored code + trtllm-gen FMHA artifacts (cubin pointers, export headers, cuda_ptx) -exclude: '(^tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_chunked_bt16(_small_bh)?\.py$|^cpp/tensorrt_llm/common/sha256/|^triton_kernels/|^tensorrt_llm/_torch/attention/backends/prims_ts/|^tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/|^tensorrt_llm/_torch/cute_dsl_kernels/cutedsl_megamoe/|^tensorrt_llm/_torch/cute_dsl_kernels/megamoe_scheduler/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/|trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo\.h$|trtllmGenKernels/gemm/trtllmGen_gemm_export/KernelMetaInfo\.h$|\.cubin\.tar\.zst$)' +exclude: '(^tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_chunked_bt16(_small_bh)?\.py$|^cpp/tensorrt_llm/common/sha256/|^triton_kernels/|^tensorrt_llm/_torch/attention/backends/prims_ts/|^tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/|^tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/|^tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/|trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo\.h$|trtllmGenKernels/gemm/trtllmGen_gemm_export/KernelMetaInfo\.h$|\.cubin\.tar\.zst$)' default_install_hook_types: [pre-commit, commit-msg] repos: diff --git a/pyproject.toml b/pyproject.toml index 6a73230989b2..d8264fc74a57 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,8 +50,8 @@ exclude = [ "tensorrt_llm/_torch/moe/kernels/rubin/rubin_contiguous_grouped_blockscaled_gemm_fused_fc12.py", "tensorrt_llm/grpc/openengine/_generated/**", # Keep vendored MegaMoE sources in upstream formatting. - "tensorrt_llm/_torch/cute_dsl_kernels/cutedsl_megamoe/**", - "tensorrt_llm/_torch/cute_dsl_kernels/megamoe_scheduler/**", + "tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/**", + "tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/**", ] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_10x__default/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_10x__default/ops/moe/__init__.py index c72d7918ac70..52a7a9daf028 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_10x__default/ops/moe/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_10x__default/ops/moe/__init__.py @@ -1,25 +1,2 @@ -"""NVLink-domain pointer mapping and token communication.""" - -from ...helpers.software_sync import NvlinkBarrier, SoftwareGridSync -from ...quant_def import QuantKind -from .symmetric_buffer import SymmetricBufferDevice, SymmetricBufferHost -from .token_comm import ( - TokenBackMode, - TokenBackScheduleMode, - TokenCommArgs, - TokenCommNonDeterministic, -) -from .token_comm_deterministic import TokenCommDeterministic - -__all__ = [ - "NvlinkBarrier", - "QuantKind", - "SoftwareGridSync", - "SymmetricBufferDevice", - "SymmetricBufferHost", - "TokenBackMode", - "TokenBackScheduleMode", - "TokenCommArgs", - "TokenCommDeterministic", - "TokenCommNonDeterministic", -] +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/attention/kernels/rubin/dsv4_qb_fusion/__init__.py b/tensorrt_llm/_torch/attention/kernels/rubin/dsv4_qb_fusion/__init__.py index cc1ef51580d3..52a7a9daf028 100644 --- a/tensorrt_llm/_torch/attention/kernels/rubin/dsv4_qb_fusion/__init__.py +++ b/tensorrt_llm/_torch/attention/kernels/rubin/dsv4_qb_fusion/__init__.py @@ -1,22 +1,2 @@ -"""Cross-rank communication protocols and implementations.""" - -from ..quant_def import CombineFormat, QuantKind -from .nvlink_domain.token_comm import ( - TokenBackScheduleMode, - TokenBackMode, - TokenCommArgs, - TokenCommNonDeterministic, -) -from .nvlink_domain.token_comm_deterministic import TokenCommDeterministic -from .token_protocol import TokenSrcMetadata - -__all__ = [ - "CombineFormat", - "QuantKind", - "TokenBackScheduleMode", - "TokenBackMode", - "TokenCommArgs", - "TokenCommDeterministic", - "TokenCommNonDeterministic", - "TokenSrcMetadata", -] +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/modules/megamoe_shared_mlp.py b/tensorrt_llm/_torch/modules/megamoe_shared_mlp.py index 34bef124376c..e041cd137e50 100644 --- a/tensorrt_llm/_torch/modules/megamoe_shared_mlp.py +++ b/tensorrt_llm/_torch/modules/megamoe_shared_mlp.py @@ -229,7 +229,7 @@ def forward( self.prepare_fc12_weights() if x.device != self._fc12_fc1_weight.device: raise ValueError("Shared FC12 input and weights must be on the same device") - from ..cute_dsl_kernels.megamoe_shared_fc12 import run_shared_fc12 + from ..moe.kernels.megamoe_shared_fc12 import run_shared_fc12 return run_shared_fc12( x, diff --git a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md index 97277d48a276..76976420ae35 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -438,9 +438,9 @@ literal also selects, descend from the CUTLASS family and stay in | `KERNEL_INTERNALS.html` | Reference for the underlying DeepGEMM kernel layout | The exported CuteDSL kernel sources for `TrtllmCutedslMegaMoeNvfp4Impl` live under -`tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/`, vendored from upstream as -recorded in `3rdparty/vendor_sources.lock.yaml`; keep local edits out of that -directory. `moe/kernels/megamoe_kernels.py` exposes the `KernelClass` entry +`tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/`, a copy of the upstream +cutedsl_megamoe sources maintained in-tree; update it by copying newer upstream +code in. `moe/kernels/megamoe_kernels.py` exposes the `KernelClass` entry points. The runner constructs `ProblemDesc` and `ImplDesc` and dispatches internally; the public backend remains `MEGAMOE_CUTEDSL`. Rubin's smaller token buckets use the genphase diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_live_arena.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_live_arena.py index 0d4566f3af44..bb1f0a025873 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_live_arena.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_live_arena.py @@ -246,8 +246,8 @@ class _Vmm: def __init__(self, device: int, group_sizes: Tuple[int, ...]) -> None: from cuda.bindings import driver as cuda - from tensorrt_llm._torch.cute_dsl_kernels.megamoe_scheduler.sami._util import check_cuda - from tensorrt_llm._torch.cute_dsl_kernels.megamoe_scheduler.sami.fabric import ( + from tensorrt_llm._torch.moe.kernels.megamoe_scheduler.sami._util import check_cuda + from tensorrt_llm._torch.moe.kernels.megamoe_scheduler.sami.fabric import ( FABRIC, _allocation_properties, _read_write_access, @@ -549,9 +549,7 @@ def __init__( comm: Any, owner_registry: List["_SharedSlotPool"], ) -> None: - from tensorrt_llm._torch.cute_dsl_kernels.megamoe_scheduler.sami import ( - hierarchy_group_sizes, - ) + from tensorrt_llm._torch.moe.kernels.megamoe_scheduler.sami import hierarchy_group_sizes torch.cuda.set_device(device) self.world = world @@ -790,7 +788,7 @@ def _flat(self, owner: _SharedLayer, pointer: int, nbytes: int) -> torch.Tensor: def build_layer(self, comm: Any) -> Tuple[Any, Dict[str, torch.Tensor], _SharedLayer, int]: """Collectively map one layer; returns its arena, framework views, owner and own allocation bytes.""" - from tensorrt_llm._torch.cute_dsl_kernels.megamoe_scheduler.sami import ( + from tensorrt_llm._torch.moe.kernels.megamoe_scheduler.sami import ( HierarchicalLiveWeightArena, LivePlaneView, LiveTerminalView, diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py index fee74b16533d..c603c55125d9 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py @@ -84,17 +84,17 @@ def __init__( ) -> None: from cuda.bindings import driver - from tensorrt_llm._torch.cute_dsl_kernels.megamoe_scheduler.cuda_scheduler import ( + from tensorrt_llm._torch.moe.kernels.megamoe_scheduler.cuda_scheduler import ( CudaPhysicalSlotScheduler, CudaSchedulerConfig, ) - from tensorrt_llm._torch.cute_dsl_kernels.megamoe_scheduler.cuda_scheduler.multirank_dist import ( + from tensorrt_llm._torch.moe.kernels.megamoe_scheduler.cuda_scheduler.multirank_dist import ( FabricSymmetricBuffer, ) - from tensorrt_llm._torch.cute_dsl_kernels.megamoe_scheduler.cuda_scheduler.runtime import ( + from tensorrt_llm._torch.moe.kernels.megamoe_scheduler.cuda_scheduler.runtime import ( symmetric_buffer_ints, ) - from tensorrt_llm._torch.cute_dsl_kernels.megamoe_scheduler.sami import ( + from tensorrt_llm._torch.moe.kernels.megamoe_scheduler.sami import ( HierarchicalSamiWeightBroadcast, ) diff --git a/tensorrt_llm/_torch/moe/fused_moe/quantization.py b/tensorrt_llm/_torch/moe/fused_moe/quantization.py index fd95c03dffc6..601de04bead5 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/moe/fused_moe/quantization.py @@ -4288,8 +4288,7 @@ def _get_fc2_alpha_input_scale( def _live_weight_bundle_layout(hidden: int, intermediate: int, expand_intermediate: Optional[int] = None): - from ...cute_dsl_kernels.megamoe_scheduler.sami.geometry import \ - BundleLayout + from ..kernels.megamoe_scheduler.sami.geometry import BundleLayout return BundleLayout.create( hidden=int(hidden), diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/__init__.py index 0e7d0e5264dc..cc1ef51580d3 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/__init__.py @@ -1,47 +1,22 @@ -"""Low-level workspace, synchronization, and PTX helpers.""" +"""Cross-rank communication protocols and implementations.""" -from .device_workspace import DeviceWorkspace -from .dsl_helpers import spin_peek, spin_wait -from .flag_batch import GpuAsyncReleaseFlagBatchTracker, GpuReleaseFlagBatchTracker, make_flag_batch_tracker -from .iket_compat import iket -from .ptx_helpers import ( - cvt_f32_to_fp8_to_f32, - cvt_f32x4_to_f8x4_pack_i32, - stg_e8m0_from_f32, - stg_e8m0x8_from_f32, -) -from .smem_workspace import SmemWorkspace -from .software_sync import NvlinkBarrier, SoftwareGridSync -from .utils import ( - IntegerType, - ceil_div, - cosize_from_shape_stride_tuples, - product, - row_major_stride, - round_up, - validate_static_integer_tuple, +from ..quant_def import CombineFormat, QuantKind +from .nvlink_domain.token_comm import ( + TokenBackScheduleMode, + TokenBackMode, + TokenCommArgs, + TokenCommNonDeterministic, ) +from .nvlink_domain.token_comm_deterministic import TokenCommDeterministic +from .token_protocol import TokenSrcMetadata __all__ = [ - "DeviceWorkspace", - "GpuAsyncReleaseFlagBatchTracker", - "GpuReleaseFlagBatchTracker", - "IntegerType", - "NvlinkBarrier", - "SmemWorkspace", - "SoftwareGridSync", - "ceil_div", - "cosize_from_shape_stride_tuples", - "cvt_f32_to_fp8_to_f32", - "cvt_f32x4_to_f8x4_pack_i32", - "make_flag_batch_tracker", - "product", - "row_major_stride", - "round_up", - "iket", - "spin_peek", - "spin_wait", - "stg_e8m0_from_f32", - "stg_e8m0x8_from_f32", - "validate_static_integer_tuple", + "CombineFormat", + "QuantKind", + "TokenBackScheduleMode", + "TokenBackMode", + "TokenCommArgs", + "TokenCommDeterministic", + "TokenCommNonDeterministic", + "TokenSrcMetadata", ] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/__init__.py index 337547020a79..c72d7918ac70 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/__init__.py @@ -1,26 +1,25 @@ -"""Kernel source modules independent of repository-only runners.""" +"""NVLink-domain pointer mapping and token communication.""" -from .schedulers import ( - BlackwellFusedFc12Scheduler, - BlockPhase, - Fc12WorkTileState, - NonSwapAbFc12WorkTileInfo, - SchedulerBase, - SchedulerConsumer, - SchedulerWorkTileBase, - SwapAbFc12WorkTileInfo, - WorkIdAcquisitionMode, +from ...helpers.software_sync import NvlinkBarrier, SoftwareGridSync +from ...quant_def import QuantKind +from .symmetric_buffer import SymmetricBufferDevice, SymmetricBufferHost +from .token_comm import ( + TokenBackMode, + TokenBackScheduleMode, + TokenCommArgs, + TokenCommNonDeterministic, ) - +from .token_comm_deterministic import TokenCommDeterministic __all__ = [ - "BlackwellFusedFc12Scheduler", - "BlockPhase", - "Fc12WorkTileState", - "NonSwapAbFc12WorkTileInfo", - "SchedulerBase", - "SchedulerConsumer", - "SchedulerWorkTileBase", - "SwapAbFc12WorkTileInfo", - "WorkIdAcquisitionMode", + "NvlinkBarrier", + "QuantKind", + "SoftwareGridSync", + "SymmetricBufferDevice", + "SymmetricBufferHost", + "TokenBackMode", + "TokenBackScheduleMode", + "TokenCommArgs", + "TokenCommDeterministic", + "TokenCommNonDeterministic", ] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/__init__.py index a84af22de952..0e7d0e5264dc 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/__init__.py @@ -1,22 +1,47 @@ -"""Blackwell kernel implementations.""" +"""Low-level workspace, synchronization, and PTX helpers.""" -from .inference import ( - BlockScaledSwapAbFc12Extension, - BlockScaledSwapAbFc12Kernel, - BlockScaledSwapAbFc12Mainloop, - BlockScaledSwapAbMegaMoeKernel, - GatedActEpilogueArgs, - SwapABGatedActEpilogue, - TensorRole, +from .device_workspace import DeviceWorkspace +from .dsl_helpers import spin_peek, spin_wait +from .flag_batch import GpuAsyncReleaseFlagBatchTracker, GpuReleaseFlagBatchTracker, make_flag_batch_tracker +from .iket_compat import iket +from .ptx_helpers import ( + cvt_f32_to_fp8_to_f32, + cvt_f32x4_to_f8x4_pack_i32, + stg_e8m0_from_f32, + stg_e8m0x8_from_f32, +) +from .smem_workspace import SmemWorkspace +from .software_sync import NvlinkBarrier, SoftwareGridSync +from .utils import ( + IntegerType, + ceil_div, + cosize_from_shape_stride_tuples, + product, + row_major_stride, + round_up, + validate_static_integer_tuple, ) - __all__ = [ - "BlockScaledSwapAbFc12Extension", - "BlockScaledSwapAbFc12Kernel", - "BlockScaledSwapAbFc12Mainloop", - "BlockScaledSwapAbMegaMoeKernel", - "GatedActEpilogueArgs", - "SwapABGatedActEpilogue", - "TensorRole", + "DeviceWorkspace", + "GpuAsyncReleaseFlagBatchTracker", + "GpuReleaseFlagBatchTracker", + "IntegerType", + "NvlinkBarrier", + "SmemWorkspace", + "SoftwareGridSync", + "ceil_div", + "cosize_from_shape_stride_tuples", + "cvt_f32_to_fp8_to_f32", + "cvt_f32x4_to_f8x4_pack_i32", + "make_flag_batch_tracker", + "product", + "row_major_stride", + "round_up", + "iket", + "spin_peek", + "spin_wait", + "stg_e8m0_from_f32", + "stg_e8m0x8_from_f32", + "validate_static_integer_tuple", ] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/__init__.py index b0d6b3eb2ace..337547020a79 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/__init__.py @@ -1,22 +1,26 @@ -"""Blackwell inference kernels.""" +"""Kernel source modules independent of repository-only runners.""" -from .mega import ( - BlockScaledSwapAbFc12Extension, - BlockScaledSwapAbFc12Kernel, - BlockScaledSwapAbFc12Mainloop, - BlockScaledSwapAbMegaMoeKernel, - GatedActEpilogueArgs, - SwapABGatedActEpilogue, - TensorRole, +from .schedulers import ( + BlackwellFusedFc12Scheduler, + BlockPhase, + Fc12WorkTileState, + NonSwapAbFc12WorkTileInfo, + SchedulerBase, + SchedulerConsumer, + SchedulerWorkTileBase, + SwapAbFc12WorkTileInfo, + WorkIdAcquisitionMode, ) __all__ = [ - "BlockScaledSwapAbFc12Extension", - "BlockScaledSwapAbFc12Kernel", - "BlockScaledSwapAbFc12Mainloop", - "BlockScaledSwapAbMegaMoeKernel", - "GatedActEpilogueArgs", - "SwapABGatedActEpilogue", - "TensorRole", + "BlackwellFusedFc12Scheduler", + "BlockPhase", + "Fc12WorkTileState", + "NonSwapAbFc12WorkTileInfo", + "SchedulerBase", + "SchedulerConsumer", + "SchedulerWorkTileBase", + "SwapAbFc12WorkTileInfo", + "WorkIdAcquisitionMode", ] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/__init__.py index 52aa7758812a..a84af22de952 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/__init__.py @@ -1,21 +1,13 @@ -"""Blackwell MegaMoE inference kernels and issue helpers.""" +"""Blackwell kernel implementations.""" -from .block_scaled_swap_ab_fc12_extension import ( +from .inference import ( BlockScaledSwapAbFc12Extension, - TensorRole, -) -from .block_scaled_swap_ab_fc12_epilogue import ( - GatedActEpilogueArgs, - SwapABGatedActEpilogue, -) -from .block_scaled_swap_ab_fc12_kernel import ( BlockScaledSwapAbFc12Kernel, -) -from .block_scaled_swap_ab_fc12_mainloop import ( BlockScaledSwapAbFc12Mainloop, -) -from .block_scaled_swap_ab_mega_moe_kernel import ( BlockScaledSwapAbMegaMoeKernel, + GatedActEpilogueArgs, + SwapABGatedActEpilogue, + TensorRole, ) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/__init__.py index b84f4d01b814..b0d6b3eb2ace 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/__init__.py @@ -1 +1,22 @@ -"""Rubin kernel implementations.""" +"""Blackwell inference kernels.""" + +from .mega import ( + BlockScaledSwapAbFc12Extension, + BlockScaledSwapAbFc12Kernel, + BlockScaledSwapAbFc12Mainloop, + BlockScaledSwapAbMegaMoeKernel, + GatedActEpilogueArgs, + SwapABGatedActEpilogue, + TensorRole, +) + + +__all__ = [ + "BlockScaledSwapAbFc12Extension", + "BlockScaledSwapAbFc12Kernel", + "BlockScaledSwapAbFc12Mainloop", + "BlockScaledSwapAbMegaMoeKernel", + "GatedActEpilogueArgs", + "SwapABGatedActEpilogue", + "TensorRole", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/__init__.py index 9f3a0e63c832..52aa7758812a 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/__init__.py @@ -1,30 +1,30 @@ -"""Rubin inference kernels.""" +"""Blackwell MegaMoE inference kernels and issue helpers.""" -from .mega import ( +from .block_scaled_swap_ab_fc12_extension import ( BlockScaledSwapAbFc12Extension, - BlockScaledSwapAbFc12Mainloop, - BlockScaledSwapAbFc12MainloopGenSpecialized, - BlockScaledSwapAbGenphaseMoeKernel, - BlockScaledSwapAbMegaMoeKernel, + TensorRole, +) +from .block_scaled_swap_ab_fc12_epilogue import ( GatedActEpilogueArgs, - GenphaseTokenComm, - GenphaseTokenCommArgs, SwapABGatedActEpilogue, - TensorRole, - dynamic_mainloop, +) +from .block_scaled_swap_ab_fc12_kernel import ( + BlockScaledSwapAbFc12Kernel, +) +from .block_scaled_swap_ab_fc12_mainloop import ( + BlockScaledSwapAbFc12Mainloop, +) +from .block_scaled_swap_ab_mega_moe_kernel import ( + BlockScaledSwapAbMegaMoeKernel, ) __all__ = [ "BlockScaledSwapAbFc12Extension", + "BlockScaledSwapAbFc12Kernel", "BlockScaledSwapAbFc12Mainloop", - "BlockScaledSwapAbFc12MainloopGenSpecialized", - "BlockScaledSwapAbGenphaseMoeKernel", "BlockScaledSwapAbMegaMoeKernel", - "GenphaseTokenComm", - "GenphaseTokenCommArgs", "GatedActEpilogueArgs", "SwapABGatedActEpilogue", "TensorRole", - "dynamic_mainloop", ] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/__init__.py index 97f3397bcf5a..b84f4d01b814 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/__init__.py @@ -1,6 +1 @@ -"""Rubin Local MegaMoE inference kernel.""" - -from .block_scaled_swap_ab_local_mega_moe_kernel import BlockScaledSwapAbLocalMegaMoeKernel - - -__all__ = ["BlockScaledSwapAbLocalMegaMoeKernel"] +"""Rubin kernel implementations.""" diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/__init__.py index 51f74b497db7..9f3a0e63c832 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/__init__.py @@ -1,13 +1,18 @@ -"""Rubin MegaMoE inference kernel components.""" +"""Rubin inference kernels.""" -from . import dynamic_mainloop -from .block_scaled_swap_ab_fc12_epilogue import GatedActEpilogueArgs, SwapABGatedActEpilogue -from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension, TensorRole -from .block_scaled_swap_ab_fc12_mainloop import BlockScaledSwapAbFc12Mainloop -from .block_scaled_swap_ab_fc12_mainloop_gen_specialized import BlockScaledSwapAbFc12MainloopGenSpecialized -from .block_scaled_swap_ab_mega_moe_kernel import BlockScaledSwapAbMegaMoeKernel -from .block_scaled_swap_ab_mega_moe_kernel_gen_specialized import BlockScaledSwapAbGenphaseMoeKernel -from .token_comm_gen_specialized import GenphaseTokenComm, GenphaseTokenCommArgs +from .mega import ( + BlockScaledSwapAbFc12Extension, + BlockScaledSwapAbFc12Mainloop, + BlockScaledSwapAbFc12MainloopGenSpecialized, + BlockScaledSwapAbGenphaseMoeKernel, + BlockScaledSwapAbMegaMoeKernel, + GatedActEpilogueArgs, + GenphaseTokenComm, + GenphaseTokenCommArgs, + SwapABGatedActEpilogue, + TensorRole, + dynamic_mainloop, +) __all__ = [ @@ -16,9 +21,9 @@ "BlockScaledSwapAbFc12MainloopGenSpecialized", "BlockScaledSwapAbGenphaseMoeKernel", "BlockScaledSwapAbMegaMoeKernel", - "GatedActEpilogueArgs", "GenphaseTokenComm", "GenphaseTokenCommArgs", + "GatedActEpilogueArgs", "SwapABGatedActEpilogue", "TensorRole", "dynamic_mainloop", diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/__init__.py index 70999789b802..97f3397bcf5a 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/__init__.py @@ -1,31 +1,6 @@ -"""Scheduler protocols and implementations.""" +"""Rubin Local MegaMoE inference kernel.""" -from .base import SchedulerBase, SchedulerConsumer, SchedulerWorkTileBase, WorkIdAcquisitionMode -from .clc_mixed_cga import ClcMixedCgaSchedulerWorker -from .fc12_mapping import ( - BlockPhase, - Fc12WorkTileState, - NonSwapAbFc12WorkTileInfo, - SwapAbFc12WorkTileInfo, - peek_ready_bit, -) -from .fc12_scheduler import BlackwellFusedFc12Scheduler, PhaseInterleavedFc12Scheduler -from .non_clc_mixed_cga import NonClcMixedCgaConfig, NonClcMixedCgaSchedulerWorker +from .block_scaled_swap_ab_local_mega_moe_kernel import BlockScaledSwapAbLocalMegaMoeKernel -__all__ = [ - "BlackwellFusedFc12Scheduler", - "BlockPhase", - "ClcMixedCgaSchedulerWorker", - "Fc12WorkTileState", - "NonSwapAbFc12WorkTileInfo", - "NonClcMixedCgaConfig", - "NonClcMixedCgaSchedulerWorker", - "PhaseInterleavedFc12Scheduler", - "SchedulerBase", - "SchedulerConsumer", - "SchedulerWorkTileBase", - "SwapAbFc12WorkTileInfo", - "WorkIdAcquisitionMode", - "peek_ready_bit", -] +__all__ = ["BlockScaledSwapAbLocalMegaMoeKernel"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/__init__.py index 52a7a9daf028..51f74b497db7 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/__init__.py @@ -1,2 +1,25 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Rubin MegaMoE inference kernel components.""" + +from . import dynamic_mainloop +from .block_scaled_swap_ab_fc12_epilogue import GatedActEpilogueArgs, SwapABGatedActEpilogue +from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension, TensorRole +from .block_scaled_swap_ab_fc12_mainloop import BlockScaledSwapAbFc12Mainloop +from .block_scaled_swap_ab_fc12_mainloop_gen_specialized import BlockScaledSwapAbFc12MainloopGenSpecialized +from .block_scaled_swap_ab_mega_moe_kernel import BlockScaledSwapAbMegaMoeKernel +from .block_scaled_swap_ab_mega_moe_kernel_gen_specialized import BlockScaledSwapAbGenphaseMoeKernel +from .token_comm_gen_specialized import GenphaseTokenComm, GenphaseTokenCommArgs + + +__all__ = [ + "BlockScaledSwapAbFc12Extension", + "BlockScaledSwapAbFc12Mainloop", + "BlockScaledSwapAbFc12MainloopGenSpecialized", + "BlockScaledSwapAbGenphaseMoeKernel", + "BlockScaledSwapAbMegaMoeKernel", + "GatedActEpilogueArgs", + "GenphaseTokenComm", + "GenphaseTokenCommArgs", + "SwapABGatedActEpilogue", + "TensorRole", + "dynamic_mainloop", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/__init__.py index 52a7a9daf028..70999789b802 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/__init__.py @@ -1,2 +1,31 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Scheduler protocols and implementations.""" + +from .base import SchedulerBase, SchedulerConsumer, SchedulerWorkTileBase, WorkIdAcquisitionMode +from .clc_mixed_cga import ClcMixedCgaSchedulerWorker +from .fc12_mapping import ( + BlockPhase, + Fc12WorkTileState, + NonSwapAbFc12WorkTileInfo, + SwapAbFc12WorkTileInfo, + peek_ready_bit, +) +from .fc12_scheduler import BlackwellFusedFc12Scheduler, PhaseInterleavedFc12Scheduler +from .non_clc_mixed_cga import NonClcMixedCgaConfig, NonClcMixedCgaSchedulerWorker + + +__all__ = [ + "BlackwellFusedFc12Scheduler", + "BlockPhase", + "ClcMixedCgaSchedulerWorker", + "Fc12WorkTileState", + "NonSwapAbFc12WorkTileInfo", + "NonClcMixedCgaConfig", + "NonClcMixedCgaSchedulerWorker", + "PhaseInterleavedFc12Scheduler", + "SchedulerBase", + "SchedulerConsumer", + "SchedulerWorkTileBase", + "SwapAbFc12WorkTileInfo", + "WorkIdAcquisitionMode", + "peek_ready_bit", +] diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_shared_fc12.py b/tensorrt_llm/_torch/moe/kernels/megamoe_shared_fc12.py index c197354ae6f6..0106651e7126 100644 --- a/tensorrt_llm/_torch/moe/kernels/megamoe_shared_fc12.py +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_shared_fc12.py @@ -7,7 +7,7 @@ import torch -from ..autotuner import AutoTuner, DynamicTensorSpec, TunableRunner, TuningConfig +from ...autotuner import AutoTuner, DynamicTensorSpec, TunableRunner, TuningConfig __all__ = [ "release_shared_fc12_cache", diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index f73117e522fa..eda6293ff984 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -48,7 +48,7 @@ l0_cpu: - unittest/_torch/kernels/kimi_k3/test_k3_mla_decode_view.py - unittest/_torch/kernels/kimi_k3/test_k3_mla_launch.py - unittest/_torch/kernels/kimi_k3/test_k3_kda_pool_layout.py - - unittest/_torch/cute_dsl_kernels/test_megamoe_shared_fc12_autotune.py + - unittest/_torch/moe/test_megamoe_shared_fc12_autotune.py - unittest/_torch/compilation/test_auto_multi_stream.py - unittest/_torch/compilation/test_phase_selective_forward.py - unittest/_torch/compilation/test_remove_copy_pass.py::test_remove_copy_preserves_minimax_producer_outputs diff --git a/tests/unittest/_torch/moe/kernels/parallel/test_quantization_sm_budget.py b/tests/unittest/_torch/kernels/parallel/test_quantization_sm_budget.py similarity index 100% rename from tests/unittest/_torch/moe/kernels/parallel/test_quantization_sm_budget.py rename to tests/unittest/_torch/kernels/parallel/test_quantization_sm_budget.py diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index 49b1bbc98f22..1cb660a908e6 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -570,7 +570,7 @@ def test_model_engine_finalizer_never_enters_collective_cleanup(): def test_executor_terminal_cleanup_is_three_phase_and_stream_scoped(monkeypatch): events = [] - module_name = "tensorrt_llm._torch.cute_dsl_kernels.megamoe_shared_fc12" + module_name = "tensorrt_llm._torch.moe.kernels.megamoe_shared_fc12" monkeypatch.setitem( sys.modules, module_name, diff --git a/tests/unittest/_torch/cute_dsl_kernels/test_megamoe_shared_fc12_autotune.py b/tests/unittest/_torch/moe/test_megamoe_shared_fc12_autotune.py similarity index 99% rename from tests/unittest/_torch/cute_dsl_kernels/test_megamoe_shared_fc12_autotune.py rename to tests/unittest/_torch/moe/test_megamoe_shared_fc12_autotune.py index c09a0e32565a..1d5b8540b45e 100644 --- a/tests/unittest/_torch/cute_dsl_kernels/test_megamoe_shared_fc12_autotune.py +++ b/tests/unittest/_torch/moe/test_megamoe_shared_fc12_autotune.py @@ -19,7 +19,7 @@ import pytest _ROOT = Path(__file__).resolve().parents[4] -_SOURCE = _ROOT / "tensorrt_llm/_torch/cute_dsl_kernels/megamoe_shared_fc12.py" +_SOURCE = _ROOT / "tensorrt_llm/_torch/moe/kernels/megamoe_shared_fc12.py" _AUTOTUNER = _ROOT / "tensorrt_llm/_torch/autotuner.py" pytestmark = pytest.mark.cpu_only From d0f2558c5783513c864cf035db05afc1549a76e1 Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 04:47:00 -0700 Subject: [PATCH 07/59] [None][chore] Sync the in-tree cutedsl_megamoe copy to upstream f0747157 Bring the in-tree MegaMoE kernel copy from upstream main 7832a330 to f0747157. Of the six upstream commits, only ab23082d (CFK-38880, Rubin combine-dedup) touches the files TensorRT-LLM carries; f2b73b8e (gen-phase tactics) and 09b6b091 (BF16 top-k reducer) change training and pre_unification paths only. - Add the opt-in combine-dedup path: TokenCommDedup, SwapABGatedActEpilogueDedup and TopkReduceDedup, selected by the impl key combine_dedup (default False). It supports the BF16 combine format only, so the default NVFP4 path keeps its kernels; the compile key gains a _nodedup suffix, which invalidates the JIT cache once. - TokenSrcMetadata packs token:u16 | topk:u8 | rank:u8 | order:u8 | count:u8 and every reader decodes through unpack(); routers reject max_tokens_per_rank above 65536. - Add token_back_mode="hybrid_warps". - Size the mixed-CGA fallback group table for every slot that can split under PDL. - Fence the Rubin FC2 generic-proxy pool stores before fc2_done. TensorRT-LLM's own edits to the copy (the SwiGLU alpha/beta port from #17956, the #19960 fc12_scheduler change, one comment and the bandit annotations) are kept. New upstream lines that compare token_* names to string literals carry the same inline "# nosec B105" as the existing ones. No TensorRT-LLM caller needs a change. Signed-off-by: jintaop --- .../communication/nvlink_domain/token_comm.py | 7 +- .../nvlink_domain/token_comm_dedup.py | 2457 +++++++++++++++++ .../nvlink_domain/token_comm_deterministic.py | 12 +- .../communication/token_protocol.py | 48 +- .../cutedsl_megamoe/helpers/ptx_helpers.py | 79 + .../block_scaled_swap_ab_fc12_epilogue.py | 18 +- ...lock_scaled_swap_ab_fc12_epilogue_dedup.py | 349 +++ .../block_scaled_swap_ab_fc12_mainloop.py | 7 +- .../block_scaled_swap_ab_mega_moe_kernel.py | 135 +- .../rubin/inference/mega/topk_reduce_dedup.py | 110 + .../schedulers/non_clc_mixed_cga.py | 11 +- 11 files changed, 3193 insertions(+), 40 deletions(-) create mode 100644 tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_dedup.py create mode 100644 tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue_dedup.py create mode 100644 tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/topk_reduce_dedup.py diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm.py index 374aa01b055b..ab052053ec9b 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm.py @@ -38,7 +38,7 @@ from .symmetric_buffer import SymmetricBufferDevice -TokenBackMode = Literal["epi_warps", "standalone_warps", "reuse_dispatch_warps"] +TokenBackMode = Literal["epi_warps", "standalone_warps", "reuse_dispatch_warps", "hybrid_warps"] TokenBackScheduleMode = Literal["static", "atomic_counter"] @@ -217,6 +217,11 @@ def _validate_router_configuration(self) -> None: value = getattr(self, field_name) if value <= 0: raise ValueError(f"{field_name} must be positive, got {value}.") + if self.max_tokens_per_rank > TokenSrcMetadata.MAX_TOKENS: + raise ValueError( + f"TokenSrcMetadata keeps the token index in {TokenSrcMetadata.TOKEN_BITS} bits; " + f"max_tokens_per_rank must be <= {TokenSrcMetadata.MAX_TOKENS}, got {self.max_tokens_per_rank}." + ) if self.expert_count % self.world_size != 0: raise ValueError( f"expert_count must be divisible by world_size, got {self.expert_count} and {self.world_size}." diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_dedup.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_dedup.py new file mode 100644 index 000000000000..8efed0023e67 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_dedup.py @@ -0,0 +1,2457 @@ +"""Combine-deduplicating token communication: dispatch as in token_comm, combine reduces a token's expert copies +on the destination rank and pushes one row per (src_rank, src_token).""" + +import dataclasses +import math +import os +from typing import Callable, ClassVar, Literal, Optional, Tuple, Union + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils +from cutlass.cute.typing import AddressSpace +from cutlass.cutlass_dsl import Int32, Int64 +from cutlass.utils.blockscaled_layout import tile_atom_to_shape_SF + +from ...api import ImplDesc, KernelComponent, OptionalRequirement, ProblemDesc +from ...helpers.device_workspace import DeviceWorkspace +from ...helpers.dsl_helpers import mark_alignment, smem_exclusive_prefix +from ...helpers.flag_batch import make_flag_batch_tracker +from ...helpers.iket_compat import iket +from ...helpers.software_sync import NvlinkBarrier +from ...helpers.ptx_helpers import ( + cp_async_bulk_s2g, + cp_reduce_async_bulk_add_bf16_s2g, + cp_reduce_async_bulk_add_u32_s2g, + nanosleep, + red_add_relaxed_sys_s32, + cvt_rn_bf16x2_f32x2, + f32x2_pack, + fma_rn_f32x2_bf16x2, + stg_b64, + stg_f32, + tma_load_1d, +) +from ...helpers.smem_workspace import SmemWorkspace +from ...helpers.utils import ceil_div, round_up +from ...quant_def import CombineFormat, QuantKind +from ..token_protocol import TokenSrcMetadata +from .symmetric_buffer import SymmetricBufferDevice + + +from .token_comm import TokenBackMode, TokenBackScheduleMode, TokenCommArgs # noqa: E402 + + + +@dataclasses.dataclass(frozen=True) +class _CopyInfo: + """Per-route order and count-minus-one, packed into two eight-bit fields.""" + + dst_order_id: Int32 + dst_token_cnt: Int32 + + FIELD_BITS: ClassVar[int] = 8 + FIELD_MASK: ClassVar[int] = (1 << FIELD_BITS) - 1 + + def pack(self) -> Int32: + return self.dst_order_id | ((self.dst_token_cnt - Int32(1)) << Int32(self.FIELD_BITS)) + + @classmethod + def unpack(cls, packed: Int32) -> "_CopyInfo": + return cls( + dst_order_id=packed & Int32(cls.FIELD_MASK), + dst_token_cnt=((packed >> Int32(cls.FIELD_BITS)) & Int32(cls.FIELD_MASK)) + Int32(1), + ) + + +@dataclasses.dataclass(frozen=True) +class _SortedElement: + """Sort-buffer element: Int128 ``[flat_topk_index | copy_info | topk_score | pad]`` with scores, else Int64 + ``[flat_topk_index | copy_info]``; ``flat_topk_index = token * topk + slot``.""" + + flat_topk_index: Int32 + copy_info: Int32 + topk_score: Optional[cutlass.Float32] + + @staticmethod + def storage_type(with_score: bool): + return cutlass.Int128 if with_score else cutlass.Int64 + + def pack(self) -> Union[cutlass.Int128, Int64]: + if cutlass.const_expr(self.topk_score is None): + scratch = cute.make_rmem_tensor((2,), cutlass.Int32) + scratch[0] = self.flat_topk_index + scratch[1] = self.copy_info + return cute.recast_tensor(scratch, cutlass.Int64)[0] + scratch = cute.make_rmem_tensor((4,), cutlass.Int32) + scratch[0] = self.flat_topk_index + scratch[1] = self.copy_info + cute.recast_tensor(scratch, cutlass.Float32)[2] = self.topk_score + return cute.recast_tensor(scratch, cutlass.Int128)[0] + + @classmethod + def from_packed(cls, packed: Union[cutlass.Int128, Int64]) -> "_SortedElement": + if cutlass.const_expr(type(packed).width == 64): + scratch = cute.make_rmem_tensor((2,), cutlass.Int32) + cute.recast_tensor(scratch, cutlass.Int64)[0] = packed + return cls(flat_topk_index=scratch[0], copy_info=scratch[1], topk_score=None) + scratch = cute.make_rmem_tensor((4,), cutlass.Int32) + cute.recast_tensor(scratch, cutlass.Int128)[0] = packed + return cls( + flat_topk_index=scratch[0], + copy_info=scratch[1], + topk_score=cute.recast_tensor(scratch, cutlass.Float32)[2], + ) + + +def _odd_even_merge_sort_pairs(length: int) -> Tuple[Tuple[int, int], ...]: + """Batcher odd-even merge sort as (low, high) compare-exchange pairs; generated for the next power of two and + pruned to pairs inside ``length`` (a pruned partner would be +inf, a no-op).""" + width = 1 << (max(length, 1) - 1).bit_length() + pairs = [] + + def merge(lo, n, r): + step = r * 2 + if step < n: + merge(lo, n, step) + merge(lo + r, n, step) + pairs.extend((i, i + r) for i in range(lo + r, lo + n - r, step)) + else: + pairs.append((lo, lo + r)) + + def sort(lo, n): + if n > 1: + half = n // 2 + sort(lo, half) + sort(lo + half, half) + merge(lo, n, 1) + + sort(0, width) + return tuple((i, j) for i, j in pairs if j < length) + + +@cute.jit +def _copy_atom(dtype, num_bits_per_copy: int): + return cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), dtype, num_bits_per_copy=num_bits_per_copy) + + +class _DedupMetadataPushRouter(KernelComponent): + """Sort and push routing metadata into each destination rank's final pool.""" + + router_smem_limit_bytes: ClassVar[int] = 227 * 1024 + router_warps_per_cta: ClassVar[int] = 16 + # Supported top-k: one token's routes are sorted in registers by a per-thread network (191 compare-exchanges at 32). + router_max_topk: ClassVar[int] = 32 + # Two Int32 registers per route (expert; within-expert index | copy info); the largest footprint without spills. + router_max_elements_per_thread: ClassVar[int] = 24 + # SMEM reserved for the expert tables / mbarriers next to the per-route buffers. + router_table_smem_bytes: ClassVar[int] = 24 * 1024 + + sizes_by_rank_region = "nvlink.token_comm_dedup.sizes_by_rank" + sizes_region = "nvlink.token_comm_dedup.sizes" + sizes_ready_region = "nvlink.token_comm_dedup.sizes_ready" + metadata_ready_region = "nvlink.token_comm_dedup.metadata_ready" + sorted_metadata_region = "nvlink.token_comm_dedup.sorted_metadata" + sorted_scores_region = "nvlink.token_comm_dedup.sorted_scores" + pool_expert_base_region = "nvlink.token_comm_dedup.pool_expert_base" + token_src_metadata_region = "nvlink.token_comm_dedup.token_src_metadata" # nosec B105 + fc1_topk_scores_region = "nvlink.token_comm_dedup.fc1_topk_scores" + source_expert_base_region = "nvlink.token_comm_dedup.source_expert_base" + push_destination_base_region = "nvlink.token_comm_dedup.push_destination_base" + sorted_metadata_ready_region = "nvlink.token_comm_dedup.sorted_metadata_ready" + push_table_ready_region = "nvlink.token_comm_dedup.push_table_ready" + router_size_counter_region = "nvlink.token_comm_dedup.router_size_counter" + router_histogram_done_region = "nvlink.token_comm_dedup.router_histogram_done" + source_base_ready_region = "nvlink.token_comm_dedup.source_base_ready" + valid_rank_mask_region = "nvlink.token_comm_dedup.valid_rank_mask" + + router_data_histogram_region = "nvlink.token_comm_dedup.router_smem.data_histogram" + router_data_prefix_region = "nvlink.token_comm_dedup.router_smem.data_prefix" + router_data_warp_totals_region = "nvlink.token_comm_dedup.router_smem.data_warp_totals" + router_data_sorted_region = "nvlink.token_comm_dedup.router_smem.data_sorted" + router_data_scores_region = "nvlink.token_comm_dedup.router_smem.data_scores" + router_data_base_region = "nvlink.token_comm_dedup.router_smem.data_base" + router_helper_size_matrix_region = "nvlink.token_comm_dedup.router_smem.helper_size_matrix" + router_helper_totals_region = "nvlink.token_comm_dedup.router_smem.helper_totals" + router_helper_prefix_region = "nvlink.token_comm_dedup.router_smem.helper_prefix" + router_push_expert_base_region = "nvlink.token_comm_dedup.router_smem.push_expert_base" + router_push_destination_base_region = "nvlink.token_comm_dedup.router_smem.push_destination_base" + router_push_warp_start_counts_region = "nvlink.token_comm_dedup.router_smem.push_warp_start_counts" + router_helper_warp_totals_region = "nvlink.token_comm_dedup.router_smem.helper_warp_totals" + router_helper_load_mbarrier_region = "nvlink.token_comm_dedup.router_smem.helper_load_mbarrier" + router_data_load_mbarrier_region = "nvlink.token_comm_dedup.router_smem.data_load_mbarrier" + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return { + "world_size": int, + "expert_count": int, + "topk": int, + "max_tokens_per_rank": int, + "apply_topk_at_fc1": bool, + } + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return { + "token_padding_block": int, + "promised_launchable_sm_count": int, + "router_smem_limit_bytes": OptionalRequirement(int), + } + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + self.world_size = problem_desc["world_size"] + self.expert_count = problem_desc["expert_count"] + self.topk = problem_desc["topk"] + self.max_tokens_per_rank = problem_desc["max_tokens_per_rank"] + self.apply_topk_at_fc1 = problem_desc["apply_topk_at_fc1"] + self.carry_topk_scores = True + + self.token_padding_block = impl_desc["token_padding_block"] + self.promised_launchable_sm_count = impl_desc["promised_launchable_sm_count"] + self.router_smem_limit_bytes = impl_desc.get("router_smem_limit_bytes", 227 * 1024) + + self._validate_router_configuration() + # Large topk: drop warps so ``threads * topk`` routes still fit the per-CTA sort capacity. + while self.router_warps_per_cta > 1 and self.router_warps_per_cta * 32 * self.topk > self.router_max_routes_per_cta: + self.router_warps_per_cta //= 2 + if 32 * self.topk > self.router_max_routes_per_cta: + raise ValueError( + f"router_smem_limit_bytes={self.router_smem_limit_bytes} leaves room for {self.router_max_routes_per_cta} " + f"routes per CTA, below the 32 x topk={32 * self.topk} that one warp holding one token per thread needs." + ) + self.expert_count_padded = round_up(self.expert_count, 4) + self.expert_count_with_trash = self.expert_count_padded + 1 + self.router_elements_per_lane, self.router_data_cta_count = self._router_launch_configuration() + # Historical name: this field counts routes (tokens_per_cta * topk) rather than source tokens. + self.router_tokens_per_cta = self.router_elements_per_lane * self.router_warps_per_cta * 32 + # Same grid footprint as the non-dedup router (DATA by route capacity, PUSH by experts / warps). + self.router_push_cta_count = min(ceil_div(self.expert_count, self.router_warps_per_cta), + self.promised_launchable_sm_count) + self.router_grid_cta_count = max(self.router_data_cta_count + 1, self.router_push_cta_count) + if self.router_grid_cta_count > self.promised_launchable_sm_count: + raise ValueError( + "Router grid exceeds promised_launchable_sm_count; all metadata-push CTAs must be concurrently resident." + ) + self.worst_case_token_count = self.worst_case_padded_tokens(self.token_padding_block) + self._router_smem_workspace = self._build_router_smem_workspace() + + self._device_workspace = None + self._peer_rank_ptr_mapper = None + self._router_local_rank = None + self._router_thread_idx = None + self._router_linear_cta_idx = None + self._router_warp_idx = None + self._router_lane_idx = None + self._router_smem_base = None + + def _validate_router_configuration(self) -> None: + if self.topk > self.router_max_topk: + raise ValueError(f"Combine dedup sorts a token's routes in registers; topk must be <= {self.router_max_topk}.") + if self.world_size > 32: + raise ValueError("Combine dedup keeps the destination-rank mask in an Int32; world_size must be <= 32.") + if self.max_tokens_per_rank > TokenSrcMetadata.MAX_TOKENS: + raise ValueError( + f"TokenSrcMetadata keeps the token index in {TokenSrcMetadata.TOKEN_BITS} bits; " + f"max_tokens_per_rank must be <= {TokenSrcMetadata.MAX_TOKENS}." + ) + positive_fields = ( + "world_size", + "expert_count", + "topk", + "max_tokens_per_rank", + "token_padding_block", + "promised_launchable_sm_count", + "router_smem_limit_bytes", + ) + for field_name in positive_fields: + value = getattr(self, field_name) + if value <= 0: + raise ValueError(f"{field_name} must be positive, got {value}.") + if self.expert_count % self.world_size != 0: + raise ValueError( + f"expert_count must be divisible by world_size, got {self.expert_count} and {self.world_size}." + ) + if self.expert_count > 16384: + raise NotImplementedError("TokenComm supports at most 16384 global experts.") + if self.topk > self.expert_count: + raise ValueError(f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}.") + + @property + def experts_per_rank(self) -> int: + return self.expert_count // self.world_size + + def worst_case_padded_tokens(self, block: int) -> int: + source_token_capacity = self.world_size * self.max_tokens_per_rank + routes_per_source_token = min(self.topk, self.experts_per_rank) + route_capacity = source_token_capacity * routes_per_source_token + active_expert_capacity = min(self.experts_per_rank, route_capacity) + route_budget_blocks = active_expert_capacity + (route_capacity - active_expert_capacity) // block + expert_bound_blocks = active_expert_capacity * int(ceil_div(source_token_capacity, block)) + return min(route_budget_blocks, expert_bound_blocks) * block + + @property + def router_max_routes_per_cta(self) -> int: + # Per-route SMEM: sort record (16 B with scores, 8 B without) + fp32 score slab; capped like the base router. + record_bytes = int(_SortedElement.storage_type(self.carry_topk_scores).width) // 8 + per_route = max(record_bytes, 8) + (4 if self.carry_topk_scores else 0) + return min(16384, (self.router_smem_limit_bytes - self.router_table_smem_bytes) // per_route) + + def _router_launch_configuration(self) -> Tuple[int, int]: + def next_power_of_two(value: int) -> int: + return 1 << (max(value, 1) - 1).bit_length() + + # DATA grid = routes / capacity; each thread owns whole tokens so a token's copies never straddle threads. + routed_token_capacity = self.max_tokens_per_rank * self.topk + threads_per_cta = self.router_warps_per_cta * 32 + minimum_cta_capacity = 2048 + cta_capacity = min( + self.router_max_routes_per_cta, next_power_of_two(max(routed_token_capacity, minimum_cta_capacity)) + ) + # Whole tokens per thread, bounded by SMEM capacity and register footprint; at least one token. + tokens_per_thread = max( + 1, min(cta_capacity // (threads_per_cta * self.topk), self.router_max_elements_per_thread // self.topk) + ) + self.router_tokens_per_thread = tokens_per_thread + if tokens_per_thread * threads_per_cta * self.topk > 1 << 16: + raise ValueError("The within-expert index is packed into 16 bits; routes per DATA CTA must be <= 65536.") + elements_per_lane = tokens_per_thread * self.topk + data_cta_count = ceil_div(self.max_tokens_per_rank, tokens_per_thread * threads_per_cta) + # Slab-read bank swizzle (ROUTER_SMEM_BANK_CONFLICT_PLAN.md section 4): lane l reads vector chunk + # ``c XOR f(l)``, f(l) = (l // (32 / gcd(K, 32))) % (gcd(K, 32) / V), which makes every wavefront cover the + # 32 banks once; ``_router_value_coordinate`` inverts it. + k = self.topk + self.router_load_vector_words = 4 if k % 4 == 0 else (2 if k % 2 == 0 else 1) + gcd_k = math.gcd(k, 32) + self.router_swizzle_period_lanes = 32 // gcd_k + self.router_swizzle_groups = max(1, gcd_k // self.router_load_vector_words) + assert (k // self.router_load_vector_words) % self.router_swizzle_groups == 0 + return elements_per_lane, data_cta_count + + @cute.jit + def _router_slot_swizzle(self) -> Int32: + """f(lane) of the slab-read swizzle: the XOR applied to a thread's vector-chunk index (0 when identity).""" + if cutlass.const_expr(self.router_swizzle_groups == 1): + return Int32(0) + lane = self._router_thread_idx % Int32(32) + return (lane // Int32(self.router_swizzle_period_lanes)) % Int32(self.router_swizzle_groups) + + def _build_router_smem_workspace(self) -> SmemWorkspace: + workspace = SmemWorkspace() + workspace.register_mbarrier(self.router_helper_load_mbarrier_region, 1) + workspace.register_mbarrier(self.router_data_load_mbarrier_region, 1) + overlay = workspace.create_overlay("nvlink.token_comm_dedup.router_smem.role") + data_lifetime = overlay.add_lifetime("data_cta") + data_lifetime.register_tensor(self.router_data_histogram_region, cutlass.Int32, (self.expert_count_with_trash,)) + data_lifetime.register_tensor( + self.router_data_prefix_region, cutlass.Int32, (self.expert_count_with_trash,), byte_alignment=16 + ) + data_lifetime.register_tensor(self.router_data_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,)) + # The sort buffer doubles as the landing slab of the bulk-copied topk indices; scores get their own slab. + sorted_element_bytes = int(_SortedElement.storage_type(self.carry_topk_scores).width) // 8 + slab_bytes = self.router_tokens_per_cta * 8 + sorted_count = max(self.router_tokens_per_cta, ceil_div(slab_bytes, sorted_element_bytes)) + data_lifetime.register_tensor( + self.router_data_sorted_region, + _SortedElement.storage_type(self.carry_topk_scores), + (sorted_count,), + byte_alignment=16, + ) + if self.carry_topk_scores: + data_lifetime.register_tensor( + self.router_data_scores_region, cutlass.Float32, (self.router_tokens_per_cta,), byte_alignment=16 + ) + if self.router_data_cta_count > 1: + data_lifetime.register_tensor( + self.router_data_base_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + + helper_lifetime = overlay.add_lifetime("helper_cta") + helper_lifetime.register_tensor( + self.router_helper_size_matrix_region, + cutlass.Int32, + (self.world_size, self.expert_count_padded), + stride=(self.expert_count_padded, 1), + byte_alignment=16, + ) + helper_lifetime.register_tensor( + self.router_helper_totals_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + helper_lifetime.register_tensor( + self.router_helper_prefix_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + helper_lifetime.register_tensor( + self.router_helper_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,), byte_alignment=16 + ) + # Dense PUSH tables: per-expert source / destination start positions plus one ballot slot per warp. + push_lifetime = overlay.add_lifetime("push_cta") + push_lifetime.register_tensor( + self.router_push_expert_base_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + push_lifetime.register_tensor( + self.router_push_destination_base_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + push_lifetime.register_tensor( + self.router_push_warp_start_counts_region, cutlass.Int32, (self.router_warps_per_cta,), byte_alignment=16 + ) + workspace.finalize(max_bytes=self.router_smem_limit_bytes) + return workspace + + @property + def router_smem_workspace(self) -> SmemWorkspace: + return self._router_smem_workspace + + def register_device_workspace(self, workspace: DeviceWorkspace) -> None: + """Register router-private state and Router-to-Main outputs.""" + self._register_router_workspace(workspace) + + def _register_router_workspace(self, workspace: DeviceWorkspace) -> None: + maximum_routed_tokens = self.max_tokens_per_rank * self.topk + workspace.register( + self.sizes_by_rank_region, + cutlass.Int32, + (self.world_size, self.expert_count_padded), + buffer_space="shared", + stride=(self.expert_count_padded, 1), + ) + workspace.register( + self.sizes_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="shared", reset="tail_reset" + ) + workspace.register(self.sizes_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.metadata_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.sorted_metadata_region, cutlass.Int64, (maximum_routed_tokens,), buffer_space="local") + if self.carry_topk_scores: + workspace.register( + self.sorted_scores_region, cutlass.Float32, (maximum_routed_tokens,), buffer_space="local" + ) + workspace.register( + self.token_src_metadata_region, + cutlass.Int64, + (self.worst_case_token_count,), + buffer_space="shared", + byte_alignment=16, + ) + if self.carry_topk_scores: + workspace.register( + self.fc1_topk_scores_region, cutlass.Float32, (self.worst_case_token_count,), buffer_space="shared" + ) + workspace.register(self.pool_expert_base_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local") + # Per source token: bitmask of destination ranks holding >= 1 copy (rewritten every launch). + workspace.register(self.valid_rank_mask_region, cutlass.Int32, (self.max_tokens_per_rank,), buffer_space="local") + workspace.register( + self.source_expert_base_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="local" + ) + workspace.register( + self.push_destination_base_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="local" + ) + workspace.register( + self.sorted_metadata_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + workspace.register(self.push_table_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset") + if self.router_data_cta_count > 1: + workspace.register( + self.router_size_counter_region, + cutlass.Int32, + (self.expert_count_with_trash,), + buffer_space="local", + reset="tail_reset", + ) + workspace.register( + self.router_histogram_done_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + workspace.register( + self.source_base_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "_DedupMetadataPushRouter": + if values: + raise ValueError("_DedupMetadataPushRouter carries no MLIR values.") + return self + + @cute.jit + def launch_router( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper_host, + device_workspace: DeviceWorkspace, + stream: cuda.CUstream, + ) -> None: + """Launch counting-sort DATA, size-exchange HELPER, and metadata PUSH roles.""" + if cutlass.const_expr(self.carry_topk_scores and topk_scores is None): + raise ValueError("Combine dedup routing requires router topk_scores.") + peer_rank_ptr_mapper = peer_rank_ptr_mapper_host.make_device_object() + self._router_kernel( + topk_indices, + topk_scores, + local_rank, + local_workspace, + shared_workspace, + peer_rank_ptr_mapper, + device_workspace, + ).launch( + grid=[self.router_grid_cta_count, 1, 1], + block=[self.router_warps_per_cta * 32, 1, 1], + min_blocks_per_mp=1, + stream=stream, + ) + + @cute.kernel + def _router_kernel( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper: SymmetricBufferDevice, + device_workspace: DeviceWorkspace, + ) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + linear_cta_idx, _, _ = cute.arch.block_idx() + cute.arch.griddepcontrol_launch_dependents() + warp_idx = cute.arch.make_warp_uniform(thread_idx // Int32(32)) + lane_idx = thread_idx % Int32(32) + + storage_type = self._router_smem_workspace.storage_class() + smem_allocator = cutlass.utils.SmemAllocator() + storage = smem_allocator.allocate(storage_type) + smem_base = storage.buffer.data_ptr() + + device_workspace.assign_device_members(local_workspace, shared_workspace) + self._device_workspace = device_workspace + self._peer_rank_ptr_mapper = peer_rank_ptr_mapper + self._router_local_rank = local_rank + self._router_thread_idx = thread_idx + self._router_linear_cta_idx = linear_cta_idx + self._router_warp_idx = warp_idx + self._router_lane_idx = lane_idx + self._router_smem_base = smem_base + + if cutlass.const_expr(self.router_data_cta_count == 1): + self._router_single_cta(topk_indices, topk_scores, smem_base) + else: + self._router_multiple_ctas(topk_indices, topk_scores, smem_base) + if linear_cta_idx < Int32(self.router_push_cta_count): + self._router_push_metadata(smem_base) + + device_workspace.remove_device_members() + self._device_workspace = None + self._peer_rank_ptr_mapper = None + self._router_local_rank = None + self._router_thread_idx = None + self._router_linear_cta_idx = None + self._router_warp_idx = None + self._router_lane_idx = None + self._router_smem_base = None + + @cute.jit + def _router_single_cta( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor], smem_base: cute.Pointer + ) -> None: + if self._router_linear_cta_idx < Int32(self.router_data_cta_count): + block_thread_count = self.router_warps_per_cta * 32 + trash_bucket = self.expert_count_padded + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) + + zero_round_count = ceil_div(self.expert_count_with_trash, block_thread_count) + for zero_round in cutlass.range_constexpr(zero_round_count): + expert = Int32(zero_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_with_trash): + histogram[expert] = Int32(0) + + iket.range_push("router.histogram") + expert_registers, score_registers = self._load_router_inputs(topk_indices, topk_scores) + cute.arch.sync_threads() + within_expert_indices = self._build_histogram(expert_registers, histogram) + iket.range_pop() + + iket.range_push("router.prefix_and_publish") + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + publish_sizes = self._broadcast_sizes_to_peers( + cute.make_tensor(histogram.iterator, cute.make_layout((self.expert_count_padded,))) + ) + total_valid_routes = smem_exclusive_prefix( + cute.make_tensor(histogram.iterator, cute.make_layout((self.expert_count_padded,))), + cute.make_tensor(prefix.iterator, cute.make_layout((self.expert_count_padded,))), + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + if self._router_thread_idx == Int32(0): + prefix[trash_bucket] = total_valid_routes + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + source_expert_base[expert] = prefix[expert] + cute.arch.sync_threads() + publish_sizes() + iket.range_pop() + + # Copy fields are computed while the helper exchanges sizes and builds push tables. + iket.range_push("router.sort") + self._sort_token_routes_and_scatter( + expert_registers, within_expert_indices, score_registers, sorted_elements, prefix, topk_indices.dtype + ) + cute.arch.sync_threads() + iket.range_pop() + iket.range_push("router.write_out") + self._dump_contiguous_router_output(sorted_elements, total_valid_routes) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" + ) + iket.range_pop() + elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): + self._router_helper_single_cta(smem_base) + + @cute.jit + def _router_multiple_ctas( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor], smem_base: cute.Pointer + ) -> None: + if self._router_linear_cta_idx < Int32(self.router_data_cta_count): + block_thread_count = self.router_warps_per_cta * 32 + trash_bucket = self.expert_count_padded + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) + dump_base = self._router_smem_workspace.tensor(self.router_data_base_region, smem_base) + + zero_round_count = ceil_div(self.expert_count_with_trash, block_thread_count) + for zero_round in cutlass.range_constexpr(zero_round_count): + expert = Int32(zero_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_with_trash): + histogram[expert] = Int32(0) + + iket.range_push("router.histogram") + expert_registers, score_registers = self._load_router_inputs(topk_indices, topk_scores) + cute.arch.sync_threads() + within_expert_indices = self._build_histogram(expert_registers, histogram) + iket.range_pop() + + iket.range_push("router.reserve_and_prefix") + size_counter = self._device_workspace.ptr(self.router_size_counter_region) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + dump_base[expert] = Int32( + cute.arch.atomic_add(size_counter + expert, histogram[expert], sem="relaxed", scope="gpu") + ) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.router_histogram_done_region), Int32(1), sem="release", scope="gpu" + ) + + total_valid_routes = smem_exclusive_prefix( + cute.make_tensor(histogram.iterator, cute.make_layout((self.expert_count_padded,))), + cute.make_tensor(prefix.iterator, cute.make_layout((self.expert_count_padded,))), + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + if self._router_thread_idx == Int32(0): + prefix[trash_bucket] = total_valid_routes + cute.arch.sync_threads() + iket.range_pop() + # Copy fields are computed while the helper exchanges sizes and builds push tables. + iket.range_push("router.sort") + self._sort_token_routes_and_scatter( + expert_registers, within_expert_indices, score_registers, sorted_elements, prefix, topk_indices.dtype + ) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.wait_source_base") + source_base_ready = self._device_workspace.ptr(self.source_base_ready_region) + if self._router_thread_idx == Int32(0): + while cute.arch.load(source_base_ready, Int32, sem="acquire", scope="gpu") != Int32(1): + nanosleep(150) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.write_out") + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + dump_base[expert] = dump_base[expert] + source_expert_base[expert] + cute.arch.sync_threads() + self._dump_router_output_by_expert(histogram, prefix, dump_base, sorted_elements) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" + ) + iket.range_pop() + elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): + self._router_helper_multiple_ctas(smem_base) + + @cute.jit + def _router_helper_single_cta(self, smem_base: cute.Pointer) -> None: + iket.range_push("router.compute_push_tables") + self._compute_push_tables(smem_base) + iket.range_pop() + + @cute.jit + def _router_helper_multiple_ctas(self, smem_base: cute.Pointer) -> None: + block_thread_count = self.router_warps_per_cta * 32 + totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) + size_counter = self._device_workspace.tensor(self.router_size_counter_region) + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + + histogram_done = self._device_workspace.ptr(self.router_histogram_done_region) + iket.range_push("router.wait_histogram") + if self._router_thread_idx == Int32(0): + while cute.arch.load(histogram_done, Int32, sem="acquire", scope="gpu") != Int32( + self.router_data_cta_count + ): + nanosleep(150) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.broadcast_sizes") + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + totals[expert] = size_counter[expert] + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + + publish_sizes = self._broadcast_sizes_to_peers(totals) + smem_exclusive_prefix( + totals, + prefix, + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + source_expert_base[expert] = prefix[expert] + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.source_base_ready_region), Int32(1), sem="release", scope="gpu" + ) + publish_sizes() + iket.range_pop() + + iket.range_push("router.compute_push_tables") + self._compute_push_tables(smem_base) + iket.range_pop() + + @cute.jit + def _router_push_metadata(self, smem_base: cute.Pointer) -> None: + """Dense PUSH of the sorted records: each CTA sweeps a contiguous interval block-stride and derives a + record's expert from its position by a forward scan over the SMEM prefix table.""" + block_thread_count = self.router_warps_per_cta * 32 + sorted_metadata_ready = self._device_workspace.ptr(self.sorted_metadata_ready_region) + push_table_ready = self._device_workspace.ptr(self.push_table_ready_region) + if self._router_thread_idx == Int32(0): + while cute.arch.load(sorted_metadata_ready, Int32, sem="acquire", scope="gpu") != Int32( + self.router_data_cta_count + ): + nanosleep(150) + while cute.arch.load(push_table_ready, Int32, sem="acquire", scope="gpu") != Int32(1): + nanosleep(150) + cute.arch.sync_threads() + + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + global_destination_base = self._device_workspace.tensor(self.push_destination_base_region) + source_sizes = self._device_workspace.tensor(self.sizes_by_rank_region) + last_expert = Int32(self.expert_count - 1) + # Invalid routes sit in the trash bucket after the last expert and are never pushed. + total_routes = source_expert_base[last_expert] + source_sizes[Int32(self._router_local_rank), last_expert] + # CTA c owns records [c * routes_per_cta, (c + 1) * routes_per_cta), a multiple of the block (coalesced). + push_threads = Int32(self.router_push_cta_count * block_thread_count) + routes_per_cta = (total_routes + push_threads - Int32(1)) // push_threads * Int32(block_thread_count) + cta_begin = self._router_linear_cta_idx * routes_per_cta + cta_end = cutlass.min(cta_begin + routes_per_cta, total_routes) + + # Stage the per-expert tables in SMEM and ballot the expert of this CTA's first record. + expert_base = self._router_smem_workspace.tensor(self.router_push_expert_base_region, smem_base) + push_destination_base = self._router_smem_workspace.tensor(self.router_push_destination_base_region, smem_base) + warp_start_counts = self._router_smem_workspace.tensor(self.router_push_warp_start_counts_region, smem_base) + warp_start_count = Int32(0) + for expert_round in cutlass.range_constexpr(ceil_div(self.expert_count_padded, block_thread_count)): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + starts_before = Int32(0) + if expert < Int32(self.expert_count_padded): + expert_start = source_expert_base[expert] + expert_base[expert] = expert_start + push_destination_base[expert] = global_destination_base[expert] + if expert <= last_expert: + if expert_start <= cta_begin: + starts_before = Int32(1) + warp_start_count = warp_start_count + Int32( + cute.arch.popc(cute.arch.vote_ballot_sync(starts_before != Int32(0))) + ) + if self._router_lane_idx == Int32(0): + warp_start_counts[self._router_warp_idx] = warp_start_count + cute.arch.sync_threads() + start_count = Int32(0) + for warp in cutlass.range_constexpr(self.router_warps_per_cta): + start_count = start_count + warp_start_counts[warp] + expert = cutlass.max(start_count - Int32(1), Int32(0)) + + source_metadata = self._device_workspace.ptr(self.sorted_metadata_region) + destination_metadata_base = self._device_workspace.ptr(self.token_src_metadata_region).toint() + if cutlass.const_expr(self.carry_topk_scores): + source_scores = self._device_workspace.ptr(self.sorted_scores_region) + destination_scores_base = self._device_workspace.ptr(self.fc1_topk_scores_region).toint() + source_position = cta_begin + self._router_thread_idx + # Per-expert state lives in registers and is refreshed only when the cursor advances. + next_expert_start = Int32(0x7FFFFFFF) + if expert < last_expert: + next_expert_start = expert_base[expert + Int32(1)] + while (expert < last_expert) and (next_expert_start <= source_position): + expert = expert + Int32(1) + if expert < last_expert: + next_expert_start = expert_base[expert + Int32(1)] + else: + next_expert_start = Int32(0x7FFFFFFF) + destination_shift = push_destination_base[expert] - expert_base[expert] + peer_offset = self._peer_rank_ptr_mapper.map(Int64(0), expert // Int32(self.experts_per_rank), Int64(0)) + while source_position < cta_end: + metadata = cute.arch.load(source_metadata + source_position, cutlass.Int64) + if next_expert_start <= source_position: + # Advance the cursor to the last expert whose first position is <= source_position. + while (expert < last_expert) and (next_expert_start <= source_position): + expert = expert + Int32(1) + if expert < last_expert: + next_expert_start = expert_base[expert + Int32(1)] + else: + next_expert_start = Int32(0x7FFFFFFF) + destination_shift = push_destination_base[expert] - expert_base[expert] + peer_offset = self._peer_rank_ptr_mapper.map(Int64(0), expert // Int32(self.experts_per_rank), Int64(0)) + destination_position = source_position + destination_shift + stg_b64(destination_metadata_base + peer_offset + Int64(destination_position) * Int64(TokenSrcMetadata.nbytes), metadata) + if cutlass.const_expr(self.carry_topk_scores): + score = cute.arch.load(source_scores + source_position, cutlass.Float32) + stg_f32(destination_scores_base + peer_offset + Int64(destination_position) * Int64(4), score) + source_position = source_position + Int32(block_thread_count) + + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_sys() + # Keep notifier threads behind the leader's system fence. + cute.arch.sync_threads() + metadata_ready_address = self._device_workspace.ptr(self.metadata_ready_region).toint() + rank_round_count = ceil_div(self.world_size, block_thread_count) + for rank_round in cutlass.range_constexpr(rank_round_count): + destination_rank = Int32(rank_round * block_thread_count) + self._router_thread_idx + if destination_rank < Int32(self.world_size): + red_add_relaxed_sys_s32( + self._peer_rank_ptr_mapper.map(metadata_ready_address, destination_rank, Int64(0)), Int32(1) + ) + + @cute.jit + def _broadcast_sizes_to_peers(self, smem_expert_counts: cute.Tensor) -> Callable[[], None]: + block_thread_count = self.router_warps_per_cta * 32 + row_bytes = Int32(self.expert_count_padded * 4) + matrix_address = self._device_workspace.ptr(self.sizes_by_rank_region).toint() + total_address = self._device_workspace.ptr(self.sizes_region).toint() + rank_round_count = ceil_div(self.world_size, self.router_warps_per_cta) + for rank_round in cutlass.range_constexpr(rank_round_count): + destination_rank = self._router_warp_idx + Int32(rank_round * self.router_warps_per_cta) + if destination_rank < Int32(self.world_size): + peer_offset = self._peer_rank_ptr_mapper.map(Int64(0), destination_rank, Int64(0)) + destination_row_address = ( + matrix_address + + peer_offset + + Int64(Int32(self._router_local_rank) * Int32(self.expert_count_padded)) * Int64(4) + ) + destination_row = cute.make_ptr( + cutlass.Int32, destination_row_address, AddressSpace.gmem, assumed_align=16 + ) + destination_total = cute.make_ptr( + cutlass.Int32, total_address + peer_offset, AddressSpace.gmem, assumed_align=16 + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_row, smem_expert_counts.iterator, row_bytes) + cp_reduce_async_bulk_add_u32_s2g(destination_total, smem_expert_counts.iterator, row_bytes) + cute.arch.cp_async_bulk_commit_group() + + def finalize() -> None: + cute.arch.cp_async_bulk_wait_group(0) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_sys() + cute.arch.sync_threads() + ready_address = self._device_workspace.ptr(self.sizes_ready_region).toint() + ready_round_count = ceil_div(self.world_size, block_thread_count) + for ready_round in cutlass.range_constexpr(ready_round_count): + destination_rank = Int32(ready_round * block_thread_count) + self._router_thread_idx + if destination_rank < Int32(self.world_size): + red_add_relaxed_sys_s32( + self._peer_rank_ptr_mapper.map(ready_address, destination_rank, Int64(0)), Int32(1) + ) + + return finalize + + @cute.jit + def _compute_push_tables(self, smem_base: cute.Pointer) -> None: + block_thread_count = self.router_warps_per_cta * 32 + owner_expert_begin = Int32(self._router_local_rank) * Int32(self.experts_per_rank) + matrix_bytes = self.world_size * self.expert_count_padded * 4 + + size_matrix = self._router_smem_workspace.tensor(self.router_helper_size_matrix_region, smem_base) + padded_totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) + load_mbarrier = self._router_smem_workspace.ptr(self.router_helper_load_mbarrier_region, smem_base) + sizes = self._device_workspace.tensor(self.sizes_region) + + if self._router_thread_idx == Int32(0): + cute.arch.mbarrier_init(load_mbarrier, 1) + + sizes_ready = self._device_workspace.ptr(self.sizes_ready_region) + iket.range_push("router.wait_sizes_ready") + if self._router_thread_idx == Int32(0): + while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="sys") != Int32(self.world_size): + nanosleep(150) + cute.arch.mbarrier_init_fence() + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.load_sizes_and_prefix") + if self._router_thread_idx == Int32(0): + cute.arch.mbarrier_arrive_and_expect_tx(load_mbarrier, Int32(matrix_bytes)) + tma_load_1d( + size_matrix.iterator, + self._device_workspace.ptr(self.sizes_by_rank_region), + load_mbarrier, + Int32(matrix_bytes), + ) + + padded_expert_rounds = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(padded_expert_rounds): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + expert_size = sizes[expert] + padded_totals[expert] = ( + (expert_size + Int32(self.token_padding_block - 1)) // Int32(self.token_padding_block) + ) * Int32(self.token_padding_block) + cute.arch.sync_threads() + smem_exclusive_prefix( + padded_totals, + prefix, + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + pool_expert_base = self._device_workspace.tensor(self.pool_expert_base_region) + local_expert_rounds = ceil_div(self.experts_per_rank, block_thread_count) + for expert_round in cutlass.range_constexpr(local_expert_rounds): + local_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if local_expert < Int32(self.experts_per_rank): + pool_expert_base[local_expert] = prefix[owner_expert_begin + local_expert] - prefix[owner_expert_begin] + + cute.arch.mbarrier_wait(load_mbarrier, 0) + iket.range_pop() + + iket.range_push("router.build_push_destinations") + push_destination_base = self._device_workspace.tensor(self.push_destination_base_region) + padded_expert_rounds = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(padded_expert_rounds): + global_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if global_expert < Int32(self.expert_count): + destination_rank = global_expert // Int32(self.experts_per_rank) + destination_expert_begin = destination_rank * Int32(self.experts_per_rank) + destination_pool_base = prefix[global_expert] - prefix[destination_expert_begin] + local_ring_position = ( + Int32(self._router_local_rank) - destination_rank + Int32(self.world_size) + ) % Int32(self.world_size) + source_ring_offset = Int32(0) + for ring_position in cutlass.range_constexpr(self.world_size): + source_rank = (destination_rank + Int32(ring_position)) % Int32(self.world_size) + if Int32(ring_position) < local_ring_position: + source_ring_offset = source_ring_offset + size_matrix[source_rank, global_expert] + push_destination_base[global_expert] = destination_pool_base + source_ring_offset + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.push_table_ready_region), Int32(1), sem="release", scope="gpu" + ) + iket.range_pop() + + @cute.jit + def _load_router_inputs( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor] + ) -> Tuple[cute.Tensor, Optional[cute.Tensor]]: + """Bulk-copy this CTA's token slab into SMEM; thread ``tid`` then owns tokens ``tid, tid + threads, ...`` + (``topk`` contiguous routes each) in registers.""" + topk = self.topk + threads_per_cta = self.router_warps_per_cta * 32 + tokens_per_thread = self.router_tokens_per_thread + tokens_per_cta = tokens_per_thread * threads_per_cta + routes_per_cta = tokens_per_cta * topk + elements_per_thread = tokens_per_thread * topk + experts_per_rank = self.experts_per_rank + sentinel_expert = Int32(self.expert_count_padded) + actual_token_count = Int32(self.max_tokens_per_rank) + index_bytes = int(topk_indices.dtype.width) // 8 + maximum_elements = self.max_tokens_per_rank * topk + + cta_token_begin = self._router_linear_cta_idx * Int32(tokens_per_cta) + cta_route_begin = cta_token_begin * Int32(topk) + valid_routes = Int32(maximum_elements) - cta_route_begin + if valid_routes > Int32(routes_per_cta): + valid_routes = Int32(routes_per_cta) + if valid_routes < Int32(0): + valid_routes = Int32(0) + + # Index slab aliases the sorted buffer; the score slab is its own region (it must survive the sort). + smem_base_address = self._router_smem_workspace.ptr(self.router_data_sorted_region, self._router_smem_base).toint() + slab_indices = cute.make_tensor( + cute.make_ptr(topk_indices.dtype, smem_base_address, AddressSpace.smem, assumed_align=16), + cute.make_layout((routes_per_cta,)), + ) + slab_scores = None + if cutlass.const_expr(self.carry_topk_scores): + slab_scores = self._router_smem_workspace.tensor(self.router_data_scores_region, self._router_smem_base) + load_mbarrier = self._router_smem_workspace.ptr(self.router_data_load_mbarrier_region, self._router_smem_base) + topk_flat = cute.make_tensor(topk_indices.iterator, cute.make_layout((maximum_elements,))) + scores_flat = None + if cutlass.const_expr(self.carry_topk_scores): + scores_flat = cute.make_tensor(topk_scores.iterator, cute.make_layout((maximum_elements,))) + + if self._router_thread_idx == Int32(0): + cute.arch.mbarrier_init(load_mbarrier, 1) + cute.arch.mbarrier_init_fence() + cute.arch.sync_threads() + + # Bulk part: the largest 16 B multiple of the valid bytes; tail: the < 16 B remainder, one element per thread. + index_bytes_valid = valid_routes * Int32(index_bytes) + index_bytes_bulk = index_bytes_valid & Int32(~15) + score_bytes_valid = valid_routes * Int32(4) + score_bytes_bulk = score_bytes_valid & Int32(~15) + expected_bytes = index_bytes_bulk + if cutlass.const_expr(self.carry_topk_scores): + expected_bytes = expected_bytes + score_bytes_bulk + if self._router_thread_idx == Int32(0): + cute.arch.mbarrier_arrive_and_expect_tx(load_mbarrier, expected_bytes) + if index_bytes_bulk > Int32(0): + tma_load_1d( + slab_indices.iterator, + cute.make_ptr( + topk_indices.dtype, + topk_flat.iterator.toint() + Int64(cta_route_begin) * Int64(index_bytes), + AddressSpace.gmem, + assumed_align=16, + ), + load_mbarrier, + index_bytes_bulk, + ) + if cutlass.const_expr(self.carry_topk_scores): + if score_bytes_bulk > Int32(0): + tma_load_1d( + slab_scores.iterator, + cute.make_ptr( + cutlass.Float32, + scores_flat.iterator.toint() + Int64(cta_route_begin) * Int64(4), + AddressSpace.gmem, + assumed_align=16, + ), + load_mbarrier, + score_bytes_bulk, + ) + index_tail_begin = index_bytes_bulk // Int32(index_bytes) + if self._router_thread_idx < valid_routes - index_tail_begin: + route = index_tail_begin + self._router_thread_idx + slab_indices[route] = topk_flat[cta_route_begin + route] + if cutlass.const_expr(self.carry_topk_scores): + score_tail_begin = score_bytes_bulk // Int32(4) + if self._router_thread_idx < valid_routes - score_tail_begin: + route = score_tail_begin + self._router_thread_idx + slab_scores[route] = scores_flat[cta_route_begin + route] + cute.arch.mbarrier_wait(load_mbarrier, 0) + cute.arch.sync_threads() + + # Token t of this thread = routes [(t * threads + tid) * topk, +topk); only expert ids go to registers. + expert_registers = cute.make_rmem_tensor((elements_per_thread,), cutlass.Int32) + slab_indices_i32 = cute.make_tensor( + cute.make_ptr(cutlass.Int32, smem_base_address, AddressSpace.smem, assumed_align=16), + cute.make_layout((routes_per_cta * (index_bytes // 4),)), + ) + # Register s holds physical slot ``(chunk(s) XOR swizzle) * V + s % V``; _router_value_coordinate inverts it. + vector_words: cutlass.Constexpr[int] = self.router_load_vector_words + swizzle = self._router_slot_swizzle() + for token_slot in cutlass.range_constexpr(tokens_per_thread): + local_token = Int32(token_slot * threads_per_cta) + self._router_thread_idx + token_idx = cta_token_begin + local_token + token_valid = token_idx < actual_token_count + slab_route_begin = local_token * Int32(topk) + for chunk in cutlass.range_constexpr(topk // vector_words): + chunk_route_begin = slab_route_begin + (Int32(chunk) ^ swizzle) * Int32(vector_words) + for word in cutlass.range_constexpr(vector_words): + register_idx = token_slot * topk + chunk * vector_words + word + expert = slab_indices_i32[(chunk_route_begin + Int32(word)) * Int32(index_bytes // 4)] + is_invalid = (cutlass.Uint32(expert) >= cutlass.Uint32(self.expert_count)) | (not token_valid) + if is_invalid: + expert = sentinel_expert + expert_registers[register_idx] = expert + + return expert_registers, slab_scores + + @cute.jit + def _sort_token_routes_and_scatter( + self, + expert_registers: cute.Tensor, + within_expert_indices: cute.Tensor, + score_slab: Optional[cute.Tensor], + sorted_elements: cute.Tensor, + expert_run_starts: cute.Tensor, + topk_index_type: type, + ) -> None: + """Sort each token's routes by expert in registers (invalid routes sink to the end); a rank's routes form one + run, so ``dst_order_id`` = position in the run and ``dst_token_cnt`` = run length. Scatter from the sorted + registers (payload = within-expert index | physical topk slot << 16).""" + topk = self.topk + threads_per_cta = self.router_warps_per_cta * 32 + tokens_per_thread = self.router_tokens_per_thread + experts_per_rank = self.experts_per_rank + expert_count = Int32(self.expert_count) + actual_token_count = Int32(self.max_tokens_per_rank) + cta_token_begin = self._router_linear_cta_idx * Int32(tokens_per_thread * threads_per_cta) + valid_rank_mask = self._device_workspace.tensor(self.valid_rank_mask_region) + pairs = _odd_even_merge_sort_pairs(topk) + for token_slot in cutlass.range_constexpr(tokens_per_thread): + local_token = Int32(token_slot * threads_per_cta) + self._router_thread_idx + token_idx = cta_token_begin + local_token + keys = cute.make_rmem_tensor((topk,), cutlass.Int32) + payload = cute.make_rmem_tensor((topk,), cutlass.Int32) + for slot in cutlass.range_constexpr(topk): + register_idx = token_slot * topk + slot + _, physical_slot = self._router_value_coordinate(register_idx, topk_index_type) + keys[slot] = expert_registers[register_idx] + payload[slot] = within_expert_indices[register_idx] | (physical_slot << Int32(16)) + for low, high in pairs: + key_low = keys[low] + key_high = keys[high] + swap_mask = Int32(0) - Int32(key_low > key_high) + keys[low] = cutlass.min(key_low, key_high) + keys[high] = cutlass.max(key_low, key_high) + payload_low = payload[low] + payload_high = payload[high] + payload_delta = (payload_low ^ payload_high) & swap_mask + payload[low] = payload_low ^ payload_delta + payload[high] = payload_high ^ payload_delta + # Forward pass: run start (first position of this rank) and the rank bit of every valid route. + order = cute.make_rmem_tensor((topk,), cutlass.Int32) + run_start = Int32(0) + run_end_expert = Int32(0) + block_mask = Int32(0) + for position in cutlass.range_constexpr(topk): + expert = keys[position] + rank = expert // Int32(experts_per_rank) + new_run = Int32(expert >= run_end_expert) + run_start = run_start + (Int32(position) - run_start) * new_run + run_end_expert = run_end_expert + ((rank + Int32(1)) * Int32(experts_per_rank) - run_end_expert) * new_run + order[position] = Int32(position) - run_start + valid = Int32(expert < expert_count) + block_mask = block_mask | ((Int32(1) << rank) & (Int32(0) - valid)) + if token_idx < actual_token_count: + valid_rank_mask[token_idx] = block_mask + # Backward pass: run length, then the record scatter in sorted order. + run_end = Int32(topk) + for position in cutlass.range_constexpr(topk - 1, -1, -1): + if cutlass.const_expr(position + 1 < topk): + starts_next = Int32(order[position + 1] == Int32(0)) + run_end = run_end + (Int32(position + 1) - run_end) * starts_next + expert = keys[position] + valid = Int32(expert < expert_count) + order_id = order[position] * valid + copy_count = Int32(1) + (run_end - (Int32(position) - order[position]) - Int32(1)) * valid + within_expert_index = payload[position] & Int32(0xFFFF) + physical_slot = payload[position] >> Int32(16) + flat_topk_index = token_idx * Int32(topk) + physical_slot + copy_info = _CopyInfo(dst_order_id=order_id, dst_token_cnt=copy_count).pack() + destination = expert_run_starts[expert] + within_expert_index + score = None + if cutlass.const_expr(self.carry_topk_scores): + score = score_slab[local_token * Int32(topk) + physical_slot] + sorted_elements[destination] = _SortedElement(flat_topk_index, copy_info, score).pack() + + @cute.jit + def _build_histogram(self, expert_registers: cute.Tensor, histogram: cute.Tensor) -> cute.Tensor: + register_count = cute.size(expert_registers) + within_expert_indices = cute.make_rmem_tensor((register_count,), cutlass.Int32) + for register_idx in cutlass.range_constexpr(register_count): + within_expert_indices[register_idx] = Int32( + cute.arch.atomic_add( + histogram.iterator + expert_registers[register_idx], Int32(1), sem="relaxed", scope="cta" + ) + ) + cute.arch.sync_threads() + return within_expert_indices + + @cute.jit + def _router_value_coordinate(self, register_idx: int, topk_index_type: type) -> Tuple[Int32, Int32]: + """(token index, topk slot) of this thread's register ``register_idx`` under the whole-token ownership.""" + threads_per_cta = self.router_warps_per_cta * 32 + tokens_per_cta = self.router_tokens_per_thread * threads_per_cta + token_slot = register_idx // self.topk + token_idx = ( + self._router_linear_cta_idx * Int32(tokens_per_cta) + + Int32(token_slot * threads_per_cta) + + self._router_thread_idx + ) + # Undo the slab-read swizzle: register slot s came from physical slot (chunk(s) XOR f(lane)) * V + s % V. + slot = register_idx % self.topk + vector_words = self.router_load_vector_words + physical_slot = (Int32(slot // vector_words) ^ self._router_slot_swizzle()) * Int32(vector_words) + Int32( + slot % vector_words + ) + return (token_idx, physical_slot) + + @cute.jit + def _route_metadata(self, flat_topk_index: Int32, copy_info: Int32) -> TokenSrcMetadata: + """Unfold a sorted element (flat route index + packed _CopyInfo) into the wire metadata record.""" + info = _CopyInfo.unpack(copy_info) + return TokenSrcMetadata( + src_rank=Int32(self._router_local_rank), + src_token=flat_topk_index // Int32(self.topk), + src_topk=flat_topk_index % Int32(self.topk), + dst_order_id=info.dst_order_id, + dst_token_cnt=info.dst_token_cnt, + ) + + @cute.jit + def _dump_contiguous_router_output(self, sorted_elements: cute.Tensor, total_valid_routes: Int32) -> None: + block_thread_count = self.router_warps_per_cta * 32 + metadata_address = self._device_workspace.ptr(self.sorted_metadata_region).toint() + if cutlass.const_expr(self.carry_topk_scores): + score_address = self._device_workspace.ptr(self.sorted_scores_region).toint() + dump_round_count = (total_valid_routes + Int32(block_thread_count - 1)) // Int32(block_thread_count) + for dump_round in cutlass.range(dump_round_count, unroll=4): + position = Int32(dump_round * block_thread_count) + self._router_thread_idx + predicate = Int32(position < total_valid_routes) + element = _SortedElement.from_packed(sorted_elements[position]) + metadata = self._route_metadata(element.flat_topk_index, element.copy_info) + stg_b64(metadata_address + Int64(position) * Int64(TokenSrcMetadata.nbytes), metadata.pack(), predicate) + if cutlass.const_expr(self.carry_topk_scores): + stg_f32(score_address + Int64(position) * Int64(4), element.topk_score, predicate) + + @cute.jit + def _dump_router_output_by_expert( + self, + histogram: cute.Tensor, + expert_run_starts: cute.Tensor, + expert_dump_bases: cute.Tensor, + sorted_elements: cute.Tensor, + ) -> None: + metadata_address = self._device_workspace.ptr(self.sorted_metadata_region).toint() + if cutlass.const_expr(self.carry_topk_scores): + score_address = self._device_workspace.ptr(self.sorted_scores_region).toint() + expert_round_count = ceil_div(self.expert_count_padded, self.router_warps_per_cta) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = self._router_warp_idx + Int32(expert_round * self.router_warps_per_cta) + if expert < Int32(self.expert_count_padded): + run_begin = expert_run_starts[expert] + run_length = histogram[expert] + dump_begin = expert_dump_bases[expert] + route_round_count = (run_length + Int32(31)) // Int32(32) + for route_round in cutlass.range(route_round_count, unroll=1): + route = Int32(route_round) * Int32(32) + self._router_lane_idx + predicate = Int32(route < run_length) + element = _SortedElement.from_packed(sorted_elements[predicate * (run_begin + route)]) + output_position = dump_begin + route + metadata = self._route_metadata(element.flat_topk_index, element.copy_info) + packed_metadata = metadata.pack() + stg_b64( + metadata_address + Int64(output_position) * Int64(TokenSrcMetadata.nbytes), + packed_metadata, + predicate, + ) + if cutlass.const_expr(self.carry_topk_scores): + stg_f32(score_address + Int64(output_position) * Int64(4), element.topk_score, predicate) + + @cute.jit + def local_expert_sizes(self, device_workspace: DeviceWorkspace, local_rank: Int32) -> cute.Tensor: + """Return this rank's contiguous expert-size view.""" + sizes = device_workspace.tensor(self.sizes_region) + expert_begin = local_rank * Int32(self.experts_per_rank) + return cute.make_tensor(sizes.iterator + expert_begin, cute.make_layout((self.experts_per_rank,))) + + @property + def metadata_ready_target(self) -> int: + return self.router_push_cta_count * self.world_size + + @cute.jit + def sizes_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.sizes_region) + + @cute.jit + def pool_expert_base_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.pool_expert_base_region) + + @cute.jit + def token_src_metadata_pointer(self, device_workspace: DeviceWorkspace) -> cute.Pointer: + return device_workspace.ptr(self.token_src_metadata_region) + + @cute.jit + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + if lane_idx == Int32(0): + sizes_ready = device_workspace.ptr(self.sizes_ready_region) + while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="sys") != Int32(self.world_size): + nanosleep(sleep_cycles) + cute.arch.sync_warp() + + @cute.jit + def wait_for_metadata_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + if lane_idx == Int32(0): + metadata_ready = device_workspace.ptr(self.metadata_ready_region) + while cute.arch.load(metadata_ready, Int32, sem="acquire", scope="sys") != Int32( + self.metadata_ready_target + ): + nanosleep(sleep_cycles) + cute.arch.sync_warp() + + @cute.jit + def token_src_metadata_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.token_src_metadata_region) + + @cute.jit + def valid_rank_mask_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.valid_rank_mask_region) + + @cute.jit + def fc1_topk_scores_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.carry_topk_scores): + return None + return device_workspace.tensor(self.fc1_topk_scores_region) + + +class TokenCommDedup(KernelComponent): + """Public fused-kernel token communication component.""" + + transfer_warp_count: ClassVar[int] = 4 + transfer_thread_count: ClassVar[int] = transfer_warp_count * 32 + # Order-buffer hidden chunk (= stage slot); 1 KB chunks measured 0.1 ms slower at DSV4. + standalone_chunk_bytes: ClassVar[int] = 2048 + reduction_round_group: ClassVar[int] = 4 + # Last-copy reduce: two copies per 1-D TMA through two buffers, so per-warp SMEM is 2 x 2 x chunk. + token_back_copies_per_fill: ClassVar[int] = 2 + token_back_fill_buffers: ClassVar[int] = 2 + # Combine-warp fc2_done prefix-wait spin cadence (ns). Diagnostic override: MEGA_DEDUP_WAIT_SLEEP_NS. + fc2_done_wait_sleep_ns: ClassVar[int] = int(os.environ.get("MEGA_DEDUP_WAIT_SLEEP_NS", "500")) + # Backoff cap (ns) of that spin; polling every 0.5 us slows the MMA/epilogue warps (DSV4 T8192: -12 us). + fc2_done_wait_sleep_max_ns: ClassVar[int] = int(os.environ.get("MEGA_DEDUP_WAIT_SLEEP_MAX_NS", "8000")) + # SMEM cap per combine warp for the fill buffers (reuse mode shares the token_in overlay). + stage_ring_budget_bytes: ClassVar[int] = 12288 + transfer_lifetime_barrier_id: ClassVar[int] = 9 + grid_sync_barrier_id: ClassVar[int] = 10 + standalone_size_barrier_id: ClassVar[int] = 11 + token_in_size_barrier_id: ClassVar[int] = 12 + + fc1_ready_region = "nvlink.token_comm_dedup.fc1_ready" + fc1_activation_region = "nvlink.token_comm_dedup.fc1_activation" + fc1_activation_sf_region = "nvlink.token_comm_dedup.fc1_activation_sf" + fc2_done_region = "nvlink.token_comm_dedup.fc2_done" + # Order buffer the FC2 epilogue scatters into at [src_token, src_rank, dst_order_id]. + fc2_order_output_region = "nvlink.token_comm_dedup.fc2_order_output" + pre_reduced_activation_region = "nvlink.token_comm_dedup.pre_reduced_activation" + pre_reduced_activation_sf_region = "nvlink.token_comm_dedup.pre_reduced_activation_sf" + token_back_schedule_region = "nvlink.token_comm_dedup.token_back_schedule" # nosec B105 + # Rank-major copies pulled by a token's first same-rank occurrence (dispatch dedup). + + token_in_mbarrier_region = "nvlink.token_comm_dedup.main_smem.token_in_mbarriers" # nosec B105 + token_back_mbarrier_region = "nvlink.token_comm_dedup.main_smem.token_back_mbarriers" # nosec B105 + expert_sizes_smem_region = "nvlink.token_comm_dedup.main_smem.expert_sizes" + token_in_activation_smem_region = "nvlink.token_comm_dedup.main_smem.token_in_activation" # nosec B105 + token_in_sf_smem_region = "nvlink.token_comm_dedup.main_smem.token_in_sf" # nosec B105 + token_back_activation_smem_region = "nvlink.token_comm_dedup.main_smem.token_back_activation" # nosec B105 + # hybrid_warps: second set of rings for the dispatch warps, overlaid on their token_in stage. + token_back_reuse_activation_smem_region = "nvlink.token_comm_dedup.main_smem.token_back_reuse_activation" # nosec B105 + token_back_sf_smem_region = "nvlink.token_comm_dedup.main_smem.token_back_sf" # nosec B105 + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return { + "world_size": int, + "expert_count": int, + "topk": int, + "max_tokens_per_rank": int, + "hidden_size": int, + "quant_kind": str, + "combine_format": CombineFormat, + "apply_topk_at_fc1": bool, + } + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return { + "token_padding_block": int, + "sf_padding_block": int, + "tokens_per_fc1_ready_slot": int, + "fc2_done_signals_per_token_tile": int, + "promised_launchable_sm_count": int, + "token_in_flag_batch": int, + "token_back_mode": str, + "token_back_schedule_mode": str, + "reduce_topk_in_kernel": bool, + "router_smem_limit_bytes": OptionalRequirement(int), + } + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + self.world_size = problem_desc["world_size"] + self.expert_count = problem_desc["expert_count"] + self.topk = problem_desc["topk"] + self.max_tokens_per_rank = problem_desc["max_tokens_per_rank"] + self.hidden_size = problem_desc["hidden_size"] + self.quant_kind = QuantKind(problem_desc["quant_kind"]) + self.combine_format = problem_desc["combine_format"] + self.apply_topk_at_fc1 = problem_desc["apply_topk_at_fc1"] + self.carry_topk_scores = True + + self.token_padding_block = impl_desc["token_padding_block"] + self.sf_padding_block = impl_desc["sf_padding_block"] + self.tokens_per_fc1_ready_slot = impl_desc["tokens_per_fc1_ready_slot"] + self.fc2_done_signals_per_token_tile = impl_desc["fc2_done_signals_per_token_tile"] + self.promised_launchable_sm_count = impl_desc["promised_launchable_sm_count"] + self.token_in_flag_batch = impl_desc["token_in_flag_batch"] + self.token_back_mode: TokenBackMode = impl_desc["token_back_mode"] + self.token_back_schedule_mode: TokenBackScheduleMode = impl_desc["token_back_schedule_mode"] + self.reduce_topk_in_kernel = impl_desc["reduce_topk_in_kernel"] + self.router_smem_limit_bytes = impl_desc.get("router_smem_limit_bytes", 227 * 1024) + self._validate_configuration() + self._router = _DedupMetadataPushRouter(problem_desc, impl_desc) + self._nvlink_barrier = NvlinkBarrier(world_size=self.world_size, barrier_id=self.grid_sync_barrier_id) + self._device_workspace = None + self._token_comm_args = None + self._local_rank = None + self._linear_cta_idx = None + self._transfer_warp_idx = None + self._lane_idx = None + self._transfer_group = None + + def _validate_configuration(self) -> None: + positive_fields = ( + "world_size", + "expert_count", + "topk", + "max_tokens_per_rank", + "hidden_size", + "token_padding_block", + "sf_padding_block", + "tokens_per_fc1_ready_slot", + "promised_launchable_sm_count", + "router_smem_limit_bytes", + ) + for field_name in positive_fields: + value = getattr(self, field_name) + if value <= 0: + raise ValueError(f"{field_name} must be positive, got {value}.") + if self.expert_count % self.world_size != 0: + raise ValueError( + f"expert_count must be divisible by world_size, got {self.expert_count} and {self.world_size}." + ) + if self.expert_count > 16384: + raise NotImplementedError("TokenComm supports at most 16384 global experts.") + if self.topk > self.expert_count: + raise ValueError(f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}.") + if self.token_back_mode not in ("standalone_warps", "reuse_dispatch_warps", "hybrid_warps"): + raise ValueError( + f"Combine dedup pushes the reduced rows from the combine warps; token_back_mode {self.token_back_mode!r} " + "is not supported (use standalone_warps, reuse_dispatch_warps or hybrid_warps)." + ) + if self.token_back_schedule_mode not in ("static", "atomic_counter"): + raise ValueError( + f"token_back_schedule_mode must be static or atomic_counter, got {self.token_back_schedule_mode!r}." + ) + if not 1 <= self.token_in_flag_batch <= 32: + raise ValueError(f"token_in_flag_batch must be in [1, 32], got {self.token_in_flag_batch}.") + if self.tokens_per_fc1_ready_slot % self.token_padding_block != 0: + raise ValueError("tokens_per_fc1_ready_slot must be divisible by token_padding_block.") + if self.token_back_enabled and self.fc2_done_signals_per_token_tile <= 0: + raise ValueError("fc2_done_signals_per_token_tile must be positive when token-back is enabled.") + if self.combine_format.is_quantized or self.combine_format.act_dtype is not cutlass.BFloat16: + raise ValueError("Combine dedup supports the BF16 combine format only.") + if self.hidden_size % 128 != 0: + raise ValueError("Combine dedup requires hidden_size divisible by the 128-wide FC2 epilogue tile.") + element_block = self.activation_sf_vector_size * 4 + if self.hidden_size % element_block != 0: + raise ValueError(f"{self.quant_kind} requires hidden_size divisible by {element_block}.") + if self.sf_padding_block % 128 != 0: + raise ValueError("sf_padding_block must be a multiple of 128.") + + @property + def experts_per_rank(self) -> int: + return self.expert_count // self.world_size + + @property + def activation_dtype(self) -> type: + return self.quant_kind.activation_dtype + + @property + def activation_sf_dtype(self) -> type: + return self.quant_kind.sf_dtype + + @property + def activation_sf_vector_size(self) -> int: + return self.quant_kind.sf_vec_size + + @property + def bytes_per_token(self) -> int: + return self.hidden_size * int(self.activation_dtype.width) // 8 + + @property + def activation_sf_hidden_padded(self) -> int: + valid_hidden = self.hidden_size // self.activation_sf_vector_size + elements_per_16_bytes = 128 // int(self.activation_sf_dtype.width) + return int(round_up(valid_hidden, elements_per_16_bytes)) + + @property + def combine_sf_hidden_padded(self) -> int: + if not self.combine_format.is_quantized: + return 0 + valid_hidden = self.hidden_size // self.combine_format.scale_block + elements_per_16_bytes = 128 // int(self.combine_format.scale_dtype.width) + return int(round_up(valid_hidden, elements_per_16_bytes)) + + @property + def token_back_push_data(self) -> bool: + # Dedup always pushes from the combine warps (reuse_dispatch_warps / standalone_warps). + return True + + @property + def output_token_bytes(self) -> int: + return self.hidden_size * int(self.combine_format.act_dtype.width) // 8 + + # The FC2 epilogue writes 128-hidden CTA tiles; a chunk must hold whole tiles so one tile never straddles two. + order_chunk_granule_bytes: ClassVar[int] = 256 + + @property + def token_back_stage_slots(self) -> int: + """Chunk slots of SMEM per combine warp: fill buffers x copies per fill.""" + return self.token_back_fill_buffers * self.token_back_copies_per_fill + + def token_back_chunk_bytes(self, available_bytes_per_warp: int) -> int: + """Bytes per hidden chunk of the order buffer / last-copy reduce: the per-warp budget divided by the + ``token_back_stage_slots`` chunk slots, capped at the standalone chunk and the row, rounded down to the 256 B + granule; the row's last chunk may be partial (padded in the buffer).""" + row_bytes = self.output_token_bytes + granule = self.order_chunk_granule_bytes + chunk = min(self.standalone_chunk_bytes, available_bytes_per_warp // self.token_back_stage_slots) + chunk = min(chunk, ceil_div(row_bytes, granule) * granule) + chunk = chunk // granule * granule + if chunk < granule: + raise ValueError( + f"Token-back fill buffers too small: {available_bytes_per_warp} B per warp for " + f"{self.token_back_stage_slots} chunk slots gives a {chunk} B chunk (need >= {granule} B)." + ) + if chunk > 32 * 16 * self.reduction_round_group: + raise ValueError(f"Token-back chunk {chunk} B exceeds one reduce round; its accumulators must fit registers.") + return chunk + + @property + def hybrid_token_back(self) -> bool: + """hybrid_warps: the standalone combine warps AND the dispatch warps (after token_in) both run token_back. + Twice the combine workers for the same chunk size; wins where the combine tail dominates (Qwen T>=8k: + -5..-7 %), loses where mainloop interference dominates (DSV4 8k: +6 %).""" + return self.token_back_mode == "hybrid_warps" # nosec B105 + + @property + def token_back_worker_groups(self) -> int: + return 2 if self.hybrid_token_back else 1 + + def _token_back_available_bytes_per_warp(self) -> int: + if self.token_back_mode in ("standalone_warps", "hybrid_warps"): + # hybrid: one chunk size for both groups (the order-buffer layout is shared). + return self.stage_ring_budget_bytes + activation_bytes = self.bytes_per_token + sf_bytes = self.activation_sf_hidden_padded * int(self.activation_sf_dtype.width) // 8 + return max(activation_bytes + sf_bytes, self.stage_ring_budget_bytes) + + @property + def order_chunk_bytes(self) -> int: + """Hidden chunk of the order buffer layout [T, R, chunks_per_row, K, chunk].""" + return self.token_back_chunk_bytes(self._token_back_available_bytes_per_warp()) + + @property + def order_chunks_per_row(self) -> int: + return ceil_div(self.output_token_bytes, self.order_chunk_bytes) + + @property + def order_block_bytes(self) -> int: + """Bytes of one (src_token, src_rank) block: chunks_per_row x K x chunk.""" + return self.order_chunks_per_row * self.topk * self.order_chunk_bytes + + @property + def token_back_push_sf(self) -> bool: + return self.combine_format.is_quantized + + @property + def token_back_enabled(self) -> bool: + return self.token_back_push_data or self.token_back_push_sf + + @property + def worst_case_token_count(self) -> int: + return self._router.worst_case_token_count + + @property + def worst_case_sf_token_count(self) -> int: + return self._router.worst_case_padded_tokens(self.sf_padding_block) + + @property + def max_fc1_ready_slot_count(self) -> int: + return self._router.worst_case_padded_tokens(self.tokens_per_fc1_ready_slot) // self.tokens_per_fc1_ready_slot + + @property + def router_smem_workspace(self) -> SmemWorkspace: + return self._router.router_smem_workspace + + @property + def expert_count_padded(self) -> int: + return self._router.expert_count_padded + + @property + def expert_count_with_trash(self) -> int: + return self._router.expert_count_with_trash + + @property + def router_elements_per_lane(self) -> int: + return self._router.router_elements_per_lane + + @property + def router_warps_per_cta(self) -> int: + return self._router.router_warps_per_cta + + @property + def router_data_cta_count(self) -> int: + return self._router.router_data_cta_count + + @property + def router_tokens_per_cta(self) -> int: + return self._router.router_tokens_per_cta + + @property + def router_push_cta_count(self) -> int: + return self._router.router_push_cta_count + + @property + def router_grid_cta_count(self) -> int: + return self._router.router_grid_cta_count + + @property + def metadata_ready_target(self) -> int: + return self._router.metadata_ready_target + + def register_device_workspace(self, workspace: DeviceWorkspace) -> None: + self._router.register_device_workspace(workspace) + self._register_main_workspace(workspace) + self._nvlink_barrier.register_device_workspace(workspace) + + @cute.jit + def launch_router( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper_host, + device_workspace: DeviceWorkspace, + stream: cuda.CUstream, + ) -> None: + self._router.launch_router( + topk_indices, + topk_scores, + local_rank, + local_workspace, + shared_workspace, + peer_rank_ptr_mapper_host, + device_workspace, + stream, + ) + + @cute.jit + def local_expert_sizes(self, device_workspace: DeviceWorkspace, local_rank: Int32) -> cute.Tensor: + return self._router.local_expert_sizes(device_workspace, local_rank) + + @cute.jit + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + self._router.wait_for_sizes_ready(device_workspace, sleep_cycles) + + @cute.jit + def token_src_metadata_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return self._router.token_src_metadata_tensor(device_workspace) + + @cute.jit + def fc1_topk_scores_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + return self._router.fc1_topk_scores_tensor(device_workspace) + + @cute.jit + def assign_device_members( + self, + *, + device_workspace: DeviceWorkspace, + token_comm_args: TokenCommArgs, + local_rank: Int32, + linear_cta_idx: Int32, + ) -> None: + self._device_workspace = device_workspace + self._token_comm_args = token_comm_args + self._local_rank = local_rank + self._linear_cta_idx = linear_cta_idx + thread_idx, _, _ = cute.arch.thread_idx() + transfer_thread_idx = thread_idx % Int32(self.transfer_thread_count) + self._transfer_warp_idx = cute.arch.make_warp_uniform(transfer_thread_idx // Int32(32)) + self._lane_idx = transfer_thread_idx % Int32(32) + # Transfer-warp group: 0 = dispatch warps (8-11), 1 = standalone combine warps (12-15). + self._transfer_group = cute.arch.make_warp_uniform( + Int32(thread_idx // Int32(self.transfer_thread_count) == Int32(3)) + ) + self._nvlink_barrier.assign_device_members(device_workspace, token_comm_args.peer_rank_ptr_mapper) + + def remove_device_members(self) -> None: + self._nvlink_barrier.remove_device_members() + self._device_workspace = None + self._token_comm_args = None + self._local_rank = None + self._linear_cta_idx = None + self._transfer_warp_idx = None + self._lane_idx = None + self._transfer_group = None + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "TokenCommDedup": + if values: + raise ValueError("TokenCommDedup carries no MLIR values.") + return self + + def _register_main_workspace(self, workspace: DeviceWorkspace) -> None: + """Register the fused kernel's GMEM regions: FC1 pools, the token-back machinery (``fc2_done`` + order + buffer) and, unless ``reduce_topk_in_kernel``, the symmetric (T, R, H) combine plane (never cleared: every + exposed cell is rewritten each launch).""" + workspace.register( + self.fc1_ready_region, + cutlass.Int32, + (self.max_fc1_ready_slot_count,), + buffer_space="local", + reset="tail_reset", + ) + activation_element_count = self.worst_case_token_count * self.hidden_size + workspace.register( + self.fc1_activation_region, + self.activation_dtype, + (activation_element_count,), + buffer_space="local", + byte_alignment=128, + ) + activation_sf_element_count = self.worst_case_sf_token_count * self.activation_sf_hidden_padded + workspace.register( + self.fc1_activation_sf_region, + self.activation_sf_dtype, + (activation_sf_element_count,), + buffer_space="local", + byte_alignment=128, + ) + + workspace.register( + self.fc2_done_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local", reset="tail_reset" + ) + # Order buffer [T, R, chunks_per_row, K, chunk]; never cleared (every byte read was written this launch). + order_bytes = self.max_tokens_per_rank * self.world_size * self.order_block_bytes + workspace.register( + self.fc2_order_output_region, + self.combine_format.act_dtype, + (order_bytes * 8 // int(self.combine_format.act_dtype.width),), + buffer_space="local", + byte_alignment=128, + ) + if self.token_back_schedule_mode == "atomic_counter": # nosec B105 + workspace.register( + self.token_back_schedule_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + + if not self.reduce_topk_in_kernel: + # One slot per pushing rank (not per topk): the receiver sums the slots named by valid_rank_mask. + workspace.register( + self.pre_reduced_activation_region, + self.combine_format.act_dtype, + (self.max_tokens_per_rank, self.world_size, self.hidden_size), + buffer_space="shared", + mem_order=(2, 1, 0), + byte_alignment=128, + ) + + def register_smem_regions(self, workspace: SmemWorkspace) -> None: + workspace.register_mbarrier(self.token_in_mbarrier_region, self.transfer_warp_count) + if self.token_back_enabled: + workspace.register_mbarrier( + self.token_back_mbarrier_region, + self.transfer_warp_count * self.token_back_worker_groups * self.token_back_fill_buffers, + ) + workspace.register_tensor( + self.expert_sizes_smem_region, cutlass.Int32, (self.experts_per_rank,), byte_alignment=16 + ) + transfer_overlay = workspace.create_overlay("nvlink.token_comm_dedup.main_smem.transfer") + token_in_lifetime = transfer_overlay.add_lifetime("token_in") + token_in_lifetime.register_tensor( + self.token_in_activation_smem_region, + self.activation_dtype, + (self.transfer_warp_count, self.hidden_size), + byte_alignment=16, + ) + token_in_lifetime.register_tensor( + self.token_in_sf_smem_region, + self.activation_sf_dtype, + (self.transfer_warp_count, (self.activation_sf_vector_size, self.activation_sf_hidden_padded)), + stride=(self.activation_sf_hidden_padded, (0, 1)), + byte_alignment=16, + ) + if not self.token_back_enabled: + return + + if self.token_back_mode in ("standalone_warps", "hybrid_warps"): + token_back_lifetime = workspace.create_overlay( + "nvlink.token_comm_dedup.main_smem.standalone_token_back" + ).add_lifetime("token_back") + else: + token_back_lifetime = transfer_overlay.add_lifetime("token_back") + + if self.token_back_push_data: + # Per warp: two fill buffers of two chunk slots; reuse mode shares the token_in overlay. + chunk_bytes = self.order_chunk_bytes + chunk_elements = chunk_bytes * 8 // int(self.combine_format.act_dtype.width) + token_back_lifetime.register_tensor( + self.token_back_activation_smem_region, + self.combine_format.act_dtype, + (self.transfer_warp_count, self.token_back_stage_slots, chunk_elements), + byte_alignment=16, + ) + if self.hybrid_token_back: + # Dispatch warps' buffers live in their own token_in overlay (free once token_in is done). + transfer_overlay.add_lifetime("token_back").register_tensor( + self.token_back_reuse_activation_smem_region, + self.combine_format.act_dtype, + (self.transfer_warp_count, self.token_back_stage_slots, chunk_elements), + byte_alignment=16, + ) + if False: # bf16-only combine: no scale stage + sf_row_bytes = self.combine_sf_hidden_padded * int(self.combine_format.scale_dtype.width) // 8 + if sf_row_bytes > available_bytes_per_warp: + raise ValueError("Token-back scale row exceeds its per-warp stage.") + token_back_lifetime.register_tensor( + self.token_back_sf_smem_region, + self.combine_format.scale_dtype, + (self.transfer_warp_count, (self.combine_format.scale_block, self.combine_sf_hidden_padded)), + stride=(self.combine_sf_hidden_padded, (0, 1)), + byte_alignment=16, + ) + + @cute.jit + def fc1_ready_counter_pointer(self, device_workspace: DeviceWorkspace) -> cute.Pointer: + return device_workspace.ptr(self.fc1_ready_region) + + @cute.jit + def fc1_activation_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return cute.make_tensor( + device_workspace.ptr(self.fc1_activation_region), + cute.make_layout((self.worst_case_token_count, self.hidden_size), stride=(self.hidden_size, 1)), + ) + + @cute.jit + def fc1_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + layout = tile_atom_to_shape_SF( + (self.worst_case_sf_token_count, self.hidden_size, 1), self.activation_sf_vector_size + ) + return cute.make_tensor(device_workspace.ptr(self.fc1_activation_sf_region), cute.select(layout, mode=[0, 1])) + + @cute.jit + def fc2_done_counter_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.token_back_enabled): + return None + return device_workspace.tensor(self.fc2_done_region) + + @cute.jit + def fc2_activation_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + """Base of the order buffer as a (T*R*K, 1, hidden) tensor. Only the base pointer and the hidden extent + (shape[2]) are meaningful: the epilogue addresses rows through the chunked layout + [T, R, chunks_per_row, K, chunk] (``dedup_chunk_bytes`` / ``dedup_chunks_per_row`` in its impl_desc).""" + rows = self.max_tokens_per_rank * self.world_size * self.topk + return cute.make_tensor( + device_workspace.ptr(self.fc2_order_output_region), + cute.make_layout((rows, 1, self.hidden_size), stride=(self.hidden_size, self.hidden_size, 1)), + ) + + @cute.jit + def valid_rank_mask_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return self._router.valid_rank_mask_tensor(device_workspace) + @cute.jit + def fc2_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + return None + @cute.jit + def pre_reduced_activation_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + """The (tokens, world_size, hidden) rank-slot staging plane, or None under in-kernel top-k reduction.""" + if cutlass.const_expr(self.reduce_topk_in_kernel): + return None + return device_workspace.tensor(self.pre_reduced_activation_region) + + @cute.jit + def pre_reduced_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + return None + @cute.jit + def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: + """Wait for pushed metadata, then pull activation payloads into local pools (same as the non-dedup dispatch: + every copy pulls its row from the source rank; dedup happens on the combine side only).""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + global_warp_idx = self._linear_cta_idx * Int32(self.transfer_warp_count) + transfer_warp_idx + global_warp_count = Int32(self.promised_launchable_sm_count * self.transfer_warp_count) + + sizes = self._router.sizes_tensor(self._device_workspace) + pool_expert_bases = self._router.pool_expert_base_tensor(self._device_workspace) + token_metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + + iket.range_push("token_in.wait_sizes_ready") + self._router.wait_for_sizes_ready(self._device_workspace) + iket.range_pop() + iket.range_push("token_in.stage_sizes") + owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) + owner_expert_begin = self._local_rank * Int32(self.experts_per_rank) + source_sizes = cute.make_tensor(sizes.iterator + owner_expert_begin, cute.make_layout((self.experts_per_rank,))) + copy_elements = 4 if self.experts_per_rank % 4 == 0 else 1 + source_size_vectors = cute.zipped_divide( + (mark_alignment(source_sizes, 16) if cutlass.const_expr(copy_elements == 4) else source_sizes), + (copy_elements,), + ) + destination_size_vectors = cute.zipped_divide(owned_sizes, (copy_elements,)) + size_vector_count = cute.size(destination_size_vectors, mode=[1]) + size_copy_atom = _copy_atom(cutlass.Int32, copy_elements * 32) + size_copy_rounds = ceil_div(size_vector_count, self.transfer_thread_count) + size_copy_registers = cute.make_rmem_tensor((copy_elements, size_copy_rounds), cutlass.Int32) + transfer_thread_idx = transfer_warp_idx * Int32(32) + lane_idx + for size_copy_round in cutlass.range_constexpr(size_copy_rounds): + vector_idx = Int32(size_copy_round * self.transfer_thread_count) + transfer_thread_idx + if vector_idx < Int32(size_vector_count): + cute.copy( + size_copy_atom, + source_size_vectors[None, vector_idx], + size_copy_registers[None, size_copy_round], + ) + iket.range_pop() + + iket.range_push("token_in.wait_metadata_ready") + self._router.wait_for_metadata_ready(self._device_workspace) + iket.range_pop() + + for size_copy_round in cutlass.range_constexpr(size_copy_rounds): + vector_idx = Int32(size_copy_round * self.transfer_thread_count) + transfer_thread_idx + if vector_idx < Int32(size_vector_count): + cute.copy( + size_copy_atom, + size_copy_registers[None, size_copy_round], + destination_size_vectors[None, vector_idx], + ) + iket.range_push("token_in.size_barrier") + token_in_size_barrier = pipeline.NamedBarrier( + barrier_id=self.token_in_size_barrier_id, num_threads=self.transfer_thread_count + ) + token_in_size_barrier.arrive_and_wait() + iket.range_pop() + if cutlass.const_expr(self.token_back_mode in ("standalone_warps", "hybrid_warps")): + sizes_ready_barrier = pipeline.NamedBarrier( + barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count + ) + sizes_ready_barrier.arrive() + + iket.range_push("token_in.pull_payload") + token_in_mbarriers = smem_workspace.ptr(self.token_in_mbarrier_region, smem_base) + token_in_activation = smem_workspace.tensor(self.token_in_activation_smem_region, smem_base) + token_in_sf = smem_workspace.tensor(self.token_in_sf_smem_region, smem_base) + warp_mbarrier = token_in_mbarriers + transfer_warp_idx + warp_activation_stage = token_in_activation[transfer_warp_idx, None] + warp_sf_stage = token_in_sf[transfer_warp_idx, (None, None)] + if lane_idx == Int32(0): + cute.arch.mbarrier_init(warp_mbarrier, 1) + cute.arch.sync_warp() + + fc1_activation_pointer = self._device_workspace.ptr(self.fc1_activation_region) + fc1_activation_sf = self.fc1_activation_sf_tensor(self._device_workspace) + fc1_ready_counter = self._device_workspace.ptr(self.fc1_ready_region) + activation_bytes = cute.cosize(warp_activation_stage) * int(self.activation_dtype.width) // 8 + activation_sf_bytes = cute.cosize(warp_sf_stage) * int(self.activation_sf_dtype.width) // 8 + sf_copy_elements = 4 + source_sf_values = cute.slice_(warp_sf_stage, (0, None)) + source_sf_vectors = cute.zipped_divide(source_sf_values, (sf_copy_elements,)) + sf_copy_atom = _copy_atom(self.activation_sf_dtype, sf_copy_elements * int(self.activation_sf_dtype.width)) + + next_dense_token = global_warp_idx + expert_valid_begin = Int32(0) + expert_sf_begin = Int32(0) + expert_ready_slot_begin = Int32(0) + pull_phase = Int32(0) + flag_tracker = make_flag_batch_tracker( + use_async=self.token_in_flag_batch == 1, + flag_address=Int64(0), + accumulated_flags=Int32(0), + phase=Int32(0), + thread_idx=lane_idx, + ) + + local_expert = Int32(0) + while local_expert < Int32(self.experts_per_rank): + expert_token_count = owned_sizes[local_expert] + expert_valid_end = expert_valid_begin + expert_token_count + pull_count = Int32(0) + if next_dense_token < expert_valid_end: + pull_count = (expert_valid_end - next_dense_token + global_warp_count - Int32(1)) // global_warp_count + + for pull_round in cutlass.range(pull_count, unroll=1): + dense_token_idx = next_dense_token + Int32(pull_round) * global_warp_count + token_in_expert = dense_token_idx - expert_valid_begin + pool_token_idx = pool_expert_bases[local_expert] + token_in_expert + sf_token_idx = expert_sf_begin + token_in_expert + ready_slot_idx = expert_ready_slot_begin + token_in_expert // Int32(self.tokens_per_fc1_ready_slot) + + metadata = TokenSrcMetadata.load( + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) + ) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), metadata.src_rank, Int64(0)) + remote_activation_address = ( + self._token_comm_args.activation.iterator.toint() + + peer_offset + + Int64(metadata.src_token) * Int64(self.bytes_per_token) + ) + remote_sf = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.activation_sf.dtype, + self._token_comm_args.activation_sf.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.activation_sf.iterator.max_alignment), + ), + self._token_comm_args.activation_sf.layout, + ) + remote_sf_row = remote_sf[Int64(metadata.src_token), None] + + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + warp_mbarrier, Int32(activation_bytes + activation_sf_bytes) + ) + tma_load_1d( + warp_activation_stage.iterator, + Int64(remote_activation_address), + warp_mbarrier, + Int32(activation_bytes), + ) + tma_load_1d( + warp_sf_stage.iterator, remote_sf_row.iterator, warp_mbarrier, Int32(activation_sf_bytes) + ) + cute.arch.sync_warp() + cute.arch.mbarrier_wait(warp_mbarrier, pull_phase) + + destination_activation = cute.make_ptr( + self.activation_dtype, + fc1_activation_pointer.toint() + Int64(pool_token_idx) * Int64(self.bytes_per_token), + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_activation, warp_activation_stage.iterator, Int32(activation_bytes)) + cute.arch.sync_warp() + cute.arch.cp_async_bulk_commit_group() + + destination_sf_row = fc1_activation_sf[Int64(sf_token_idx), ((None, None), None)] + destination_sf_values = cute.slice_(destination_sf_row, (0, None, None)) + destination_sf_values = cute.group_modes(destination_sf_values, 0, 2) + destination_sf_vectors = cute.zipped_divide(destination_sf_values, (sf_copy_elements,)) + sf_vector_count = cute.size(destination_sf_vectors, mode=[1]) + for sf_round in cutlass.range_constexpr(ceil_div(sf_vector_count, 32)): + sf_vector_idx = Int32(sf_round * 32) + lane_idx + if sf_vector_idx < Int32(sf_vector_count): + cute.copy( + sf_copy_atom, + source_sf_vectors[None, sf_vector_idx], + destination_sf_vectors[None, sf_vector_idx], + ) + + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.sync_warp() + ready_address = (fc1_ready_counter + ready_slot_idx).toint() + flag_tracker = flag_tracker.accumulate(Int32(0), self.token_in_flag_batch, ready_address) + cute.arch.sync_warp() + pull_phase = pull_phase ^ Int32(1) + + next_dense_token = next_dense_token + pull_count * global_warp_count + expert_valid_begin = expert_valid_end + expert_sf_begin = expert_sf_begin + ( + (expert_token_count + Int32(self.sf_padding_block - 1)) // Int32(self.sf_padding_block) + ) * Int32(self.sf_padding_block) + expert_ready_slot_begin = expert_ready_slot_begin + ( + (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) + // Int32(self.tokens_per_fc1_ready_slot) + ) + local_expert = local_expert + Int32(1) + + flag_tracker.fire() + cute.arch.sync_warp() + iket.range_pop() + if cutlass.const_expr( + self.token_back_enabled and self.token_back_mode in ("reuse_dispatch_warps", "hybrid_warps") + ): + iket.range_push("token_in.transfer_barrier") + transfer_lifetime_barrier = pipeline.NamedBarrier( + barrier_id=self.transfer_lifetime_barrier_id, num_threads=self.transfer_thread_count + ) + transfer_lifetime_barrier.arrive_and_wait() + iket.range_pop() + + @cute.jit + def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: + """Combine warps: walk the pool metadata in expert order, wait for the ``fc2_done`` prefix up to the current + expert, and for every last copy (``order == cnt - 1``) reduce the block's copies in fp32 (order 0..cnt-1), + round once to bf16 and push it to the source rank. Single-copy rows are pushed by the FC2 epilogue.""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + if cutlass.const_expr(not self.token_back_enabled): + return + + workers_per_cta = self.transfer_warp_count * self.token_back_worker_groups + worker_in_cta = transfer_warp_idx + if cutlass.const_expr(self.hybrid_token_back): + worker_in_cta = self._transfer_group * Int32(self.transfer_warp_count) + transfer_warp_idx + global_worker_idx = self._linear_cta_idx * Int32(workers_per_cta) + worker_in_cta + global_worker_count = Int32(self.promised_launchable_sm_count * workers_per_cta) + token_back_mbarriers = smem_workspace.ptr(self.token_back_mbarrier_region, smem_base) + worker_mbarriers = token_back_mbarriers + worker_in_cta * Int32(self.token_back_fill_buffers) + token_back_stages = smem_workspace.tensor(self.token_back_activation_smem_region, smem_base) + worker_stages = token_back_stages[transfer_warp_idx, None, None] # (stage_slots, chunk_elements) + if cutlass.const_expr(self.hybrid_token_back): + # Dispatch warps (group 0) reduce in the rings overlaid on their finished token_in stage. + reuse_stages = smem_workspace.tensor(self.token_back_reuse_activation_smem_region, smem_base) + if self._transfer_group == Int32(0): + worker_stages = reuse_stages[transfer_warp_idx, None, None] + chunk_bytes = cute.size(worker_stages, mode=[1]) * int(self.combine_format.act_dtype.width) // 8 + if lane_idx == Int32(0): + for buffer_idx in cutlass.range_constexpr(self.token_back_fill_buffers): + cute.arch.mbarrier_init(worker_mbarriers + buffer_idx, 1) + cute.arch.sync_warp() + + owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) + if cutlass.const_expr(self.token_back_mode in ("standalone_warps", "hybrid_warps")): + # The dispatch warps already met this barrier inside token_in (hybrid); only the standalone group waits. + if cutlass.const_expr(self.token_back_mode == "standalone_warps") or self._transfer_group == Int32(1): # nosec B105 + sizes_ready_barrier = pipeline.NamedBarrier( + barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count + ) + sizes_ready_barrier.arrive_and_wait() + + pool_expert_bases = self._router.pool_expert_base_tensor(self._device_workspace) + token_metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + fc2_done = self._device_workspace.ptr(self.fc2_done_region) + order_pointer = self._device_workspace.ptr(self.fc2_order_output_region) + output_token_bytes = self.output_token_bytes + # Fills issued by this worker so far: fill s uses buffer s % 2 and phase (s // 2) & 1 of its mbarrier. + fill_sequence = Int32(0) + + iket.range_push("token_back.work") + next_dense_token = global_worker_idx - global_worker_count + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 + next_dense_token = Int32(0) + next_dense_token = self.next_token(next_dense_token) + completed_expert_prefix = Int32(0) + expert_valid_begin = Int32(0) + local_expert = Int32(0) + while local_expert < Int32(self.experts_per_rank): + expert_token_count = owned_sizes[local_expert] + expert_valid_end = expert_valid_begin + expert_token_count + while next_dense_token < expert_valid_end: + token_in_expert = next_dense_token - expert_valid_begin + pool_token_idx = pool_expert_bases[local_expert] + token_in_expert + metadata = TokenSrcMetadata.load( + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) + ) + # Single-copy rows were pushed straight from the FC2 epilogue; only multi-copy blocks are reduced here. + if (metadata.is_last_copy != Int32(0)) and (metadata.dst_token_cnt > Int32(1)): + iket.range_push("token_back.wait_fc2_warp_prefix") + wait_sleep_ns = Int32(self.fc2_done_wait_sleep_ns) + while completed_expert_prefix <= local_expert: + prefix_base = (completed_expert_prefix // Int32(32)) * Int32(32) + checked_expert = prefix_base + Int32(31) - lane_idx + lane_ready = checked_expert >= Int32(self.experts_per_rank) + if checked_expert < Int32(self.experts_per_rank): + checked_count = owned_sizes[checked_expert] + checked_target = ((checked_count + Int32(self.tokens_per_fc1_ready_slot - 1)) + // Int32(self.tokens_per_fc1_ready_slot)) * Int32(self.fc2_done_signals_per_token_tile) + lane_ready = cute.arch.load(fc2_done + checked_expert, Int32, + sem="acquire", scope="gpu") >= checked_target + ready_mask = cute.arch.vote_ballot_sync(lane_ready) + completed_expert_prefix = cutlass.min( + prefix_base + Int32(cute.arch.clz(~ready_mask)), Int32(self.experts_per_rank)) + if completed_expert_prefix <= local_expert: + nanosleep(wait_sleep_ns) + wait_sleep_ns = cutlass.min(wait_sleep_ns * Int32(2), Int32(self.fc2_done_wait_sleep_max_ns)) + # Order every lane's fc2_done acquire before the TMA that reads the completed copies. + cute.arch.sync_warp() + iket.range_pop() + block_index = Int64(metadata.src_token) * Int64(self.world_size) + Int64(metadata.src_rank) + block_address = order_pointer.toint() + block_index * Int64(self.order_block_bytes) + fill_sequence = self._reduce_and_push( + block_address, + metadata.dst_token_cnt, + metadata.src_rank, + metadata.src_token, + worker_mbarriers, + worker_stages, + fill_sequence, + chunk_bytes, + ) + next_dense_token = self.next_token(next_dense_token) + + expert_valid_begin = expert_valid_end + local_expert = local_expert + Int32(1) + iket.range_pop() + + @cute.jit + def _issue_fill( + self, + sequence: Int32, + fill_idx: Int32, + fills_per_chunk: Int32, + copy_count: Int32, + worker_mbarriers: cute.Pointer, + stage_base, + block_address: Int64, + chunk_bytes: int, + ) -> None: + """Elected lane: one 1-D TMA of the two copies of fill ``fill_idx`` (chunk-major, copies ascending) into + buffer ``sequence % 2``, arming that buffer's mbarrier with the byte count.""" + chunk_idx = fill_idx // fills_per_chunk + first_copy = (fill_idx - chunk_idx * fills_per_chunk) * Int32(self.token_back_copies_per_fill) + copies = cutlass.min(Int32(self.token_back_copies_per_fill), copy_count - first_copy) + buffer_idx = sequence % Int32(self.token_back_fill_buffers) + mbarrier = worker_mbarriers + buffer_idx + destination = cute.make_ptr( + cutlass.Uint8, + stage_base + buffer_idx * Int32(self.token_back_copies_per_fill * chunk_bytes), + AddressSpace.smem, + assumed_align=16, + ) + source = block_address + (Int64(chunk_idx) * Int64(self.topk) + Int64(first_copy)) * Int64(chunk_bytes) + cute.arch.mbarrier_arrive_and_expect_tx(mbarrier, copies * Int32(chunk_bytes)) + tma_load_1d(destination, source, mbarrier, copies * Int32(chunk_bytes)) + + @cute.jit + def _reduce_and_push( + self, + block_address: Int64, + copy_count: Int32, + src_rank: Int32, + src_token: Int32, + worker_mbarriers: cute.Pointer, + worker_stages: cute.Tensor, + fill_sequence: Int32, + chunk_bytes: int, + ) -> Int32: + """Reduce the ``copy_count`` copies of the order block at ``block_address`` chunk by chunk (two copies per + TMA through two buffers, fp32 accumulation in order 0..cnt-1, one bf16 rounding) and push each chunk. + Returns the advanced fill sequence.""" + lane_idx = self._lane_idx + copies_per_fill = self.token_back_copies_per_fill + output_token_bytes = self.output_token_bytes + chunk_count = ceil_div(output_token_bytes, chunk_bytes) + round_group = min(self.reduction_round_group, ceil_div(chunk_bytes // 16, 32)) + destination_slot = Int32(self._local_rank) + if cutlass.const_expr(self.reduce_topk_in_kernel): + destination_slot = Int32(0) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), src_rank, Int64(0)) + remote_activation = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.pre_reduced_activation.dtype, + self._token_comm_args.pre_reduced_activation.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.pre_reduced_activation.iterator.max_alignment), + ), + self._token_comm_args.pre_reduced_activation.layout, + ) + destination_row = remote_activation[Int64(src_token), destination_slot, None] + word_atom = _copy_atom(cutlass.Int32, 128) + vector_layout = cute.make_layout((4,)) + stage_base = worker_stages.iterator.toint() + one2 = f32x2_pack(cutlass.Float32(1.0), cutlass.Float32(1.0)) # {1.0, 1.0} weight for the FFMA2 adds + fills_per_chunk = (copy_count + Int32(copies_per_fill - 1)) // Int32(copies_per_fill) + fill_count = fills_per_chunk * Int32(chunk_count) + + iket.range_push("token_back.push_data") + with cute.arch.elect_one(): + self._issue_fill(fill_sequence, Int32(0), fills_per_chunk, copy_count, worker_mbarriers, stage_base, block_address, chunk_bytes) + # Whole-chunk fp32 accumulators (round_group x 32 lanes x 16 B covers the chunk, checked at construction). + words = cute.make_rmem_tensor((4,), cutlass.Int32) + acc = cute.make_rmem_tensor((4, round_group), cutlass.Int64) + # Rolled loop over the fills of the token: the body exists once in SASS (icache footprint). + for fill_idx in cutlass.range(fill_count, unroll=1): + chunk_idx = fill_idx // fills_per_chunk + fill_in_chunk = fill_idx - chunk_idx * fills_per_chunk + first_copy = fill_in_chunk * Int32(copies_per_fill) + sequence = fill_sequence + fill_idx + buffer_base = stage_base + (sequence % Int32(self.token_back_fill_buffers)) * Int32(copies_per_fill * chunk_bytes) + current_vectors = cutlass.min(Int32(chunk_bytes), Int32(output_token_bytes) - chunk_idx * Int32(chunk_bytes)) // Int32(16) + # Prefetch fill i+1 into the other buffer. If that buffer just fed a bulk copy, wait for its READ only. + if fill_idx + Int32(1) < fill_count: + cute.arch.cp_async_bulk_wait_group(0, read=True) + with cute.arch.elect_one(): + self._issue_fill(sequence + Int32(1), fill_idx + Int32(1), fills_per_chunk, copy_count, + worker_mbarriers, stage_base, block_address, chunk_bytes) + if fill_in_chunk == Int32(0): + for group in cutlass.range_constexpr(round_group): + for word in cutlass.range_constexpr(4): + acc[word, group] = Int64(0) + cute.arch.mbarrier_wait( + worker_mbarriers + sequence % Int32(self.token_back_fill_buffers), + (sequence // Int32(self.token_back_fill_buffers)) & Int32(1), + ) + iket.range_push("token_back.reduce") + # Copies accumulate in 0..cnt-1 order; groups are 512 B apart so adjacent lanes keep contiguous 16 B. + for copy_in_fill in cutlass.range_constexpr(copies_per_fill): + if first_copy + Int32(copy_in_fill) < copy_count: + for group in cutlass.range_constexpr(round_group): + vector_idx = Int32(group * 32) + lane_idx + if vector_idx < current_vectors: + slot_vector = cute.make_tensor( + cute.make_ptr(cutlass.Int32, buffer_base + Int32(copy_in_fill * chunk_bytes) + vector_idx * Int32(16), + AddressSpace.smem, assumed_align=16), + vector_layout, + ) + cute.copy(word_atom, slot_vector, words) + for word in cutlass.range_constexpr(4): + acc[word, group] = fma_rn_f32x2_bf16x2(words[word], one2, acc[word, group]) + iket.range_pop() + if fill_in_chunk == fills_per_chunk - Int32(1): + # Round once to bf16 into slot 0 of this buffer and push the chunk from there. + for group in cutlass.range_constexpr(round_group): + vector_idx = Int32(group * 32) + lane_idx + if vector_idx < current_vectors: + slot0_vector = cute.make_tensor( + cute.make_ptr(cutlass.Int32, buffer_base + vector_idx * Int32(16), AddressSpace.smem, assumed_align=16), + vector_layout, + ) + for word in cutlass.range_constexpr(4): + words[word] = cvt_rn_bf16x2_f32x2(acc[word, group]) + cute.copy(word_atom, words, slot0_vector) + # Generic-proxy SMEM writes must be visible to the bulk copy (async proxy) that reads slot 0. + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_warp() + chunk_source = cute.make_ptr(cutlass.Uint8, buffer_base, AddressSpace.smem, assumed_align=16) + destination_chunk = cute.make_ptr( + cutlass.Uint8, + destination_row.iterator.toint() + Int64(chunk_idx) * Int64(chunk_bytes), + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + if cutlass.const_expr(self.reduce_topk_in_kernel): + cp_reduce_async_bulk_add_bf16_s2g(destination_chunk, chunk_source, current_vectors * Int32(16)) + else: + cp_async_bulk_s2g(destination_chunk, chunk_source, current_vectors * Int32(16)) + cute.arch.cp_async_bulk_commit_group() + cute.arch.sync_warp() + # The pushes of this token must have landed before the worker moves on (matches the previous per-chunk wait). + cute.arch.cp_async_bulk_wait_group(0) + iket.range_pop() + cute.arch.sync_warp() + return fill_sequence + fill_count + + @cute.jit + def next_token(self, current_token: Int32) -> Int32: + # Static stride over every token_back worker in the grid; must match ``global_worker_count`` in ``token_back``. + workers_per_cta = self.transfer_warp_count * self.token_back_worker_groups + global_worker_count = self.promised_launchable_sm_count * workers_per_cta + schedule_counter = None + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 + schedule_counter = self._device_workspace.ptr(self.token_back_schedule_region) + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 + claimed_token = Int32(0) + if self._lane_idx == Int32(0): + claimed_token = cute.arch.atomic_add(schedule_counter, Int32(1), sem="relaxed", scope="gpu") + return Int32(cute.arch.shuffle_sync(claimed_token, Int32(0))) + return current_token + global_worker_count + + @cute.jit + def reset_tail(self) -> None: + """Reset communication state with the four token-in transfer warps.""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + transfer_thread_idx = transfer_warp_idx * Int32(32) + lane_idx + iket.range_push("tail.nvlink_drain") + self._nvlink_barrier.arrive_and_wait( + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + prologue_grid_sync=True, + epilogue_grid_sync=False, + ) + iket.range_pop() + total_reset_threads = self.promised_launchable_sm_count * self.transfer_thread_count + global_reset_thread = self._linear_cta_idx * Int32(self.transfer_thread_count) + transfer_thread_idx + iket.range_push("tail.reset_workspace") + self._device_workspace.reset_tail_space("shared", global_reset_thread, total_reset_threads) + self._device_workspace.reset_tail_space("local", global_reset_thread, total_reset_threads) + iket.range_pop() + iket.range_push("tail.nvlink_publish") + self._nvlink_barrier.arrive_and_wait( + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + prologue_grid_sync=True, + epilogue_grid_sync=False, + ) + iket.range_pop() + if cutlass.const_expr(os.environ.get("MEGA_USE_NCU", "0") == "1"): + iket.range_push("tail.ncu_finalize") + self._nvlink_barrier.finalize( + 2, + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + ) + iket.range_pop() + + +__all__ = ["TokenCommDedup"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_deterministic.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_deterministic.py index 6838edb87603..98b651a96143 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_deterministic.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_deterministic.py @@ -274,6 +274,11 @@ def _validate_router_configuration(self) -> None: value = getattr(self, field_name) if value <= 0: raise ValueError(f"{field_name} must be positive, got {value}.") + if self.max_tokens_per_rank > TokenSrcMetadata.MAX_TOKENS: + raise ValueError( + f"TokenSrcMetadata keeps the token index in {TokenSrcMetadata.TOKEN_BITS} bits; " + f"max_tokens_per_rank must be <= {TokenSrcMetadata.MAX_TOKENS}, got {self.max_tokens_per_rank}." + ) if self.expert_count % self.world_size != 0: raise ValueError( f"expert_count must be divisible by world_size, got {self.expert_count} and {self.world_size}." @@ -833,8 +838,11 @@ def _router_push_metadata(self) -> None: physical_store, ) if cutlass.const_expr(self.input_dedup_mode != "off"): - source_token = Int32(metadata & Int64(0xFFFFFFFF)) - source_slot = Int32((metadata >> Int64(32)) & Int64(0xFFFF)) + # Decode through the protocol class: the record layout changed (token:u16 | topk:u8 | + # rank:u8 | order:u8 | cnt:u8) and a hand-rolled mask here silently indexed the wrong route. + decoded = TokenSrcMetadata.unpack(Int64(metadata)) + source_token = decoded.src_token + source_slot = decoded.src_topk route_index = (self._router_local_rank * Int32(self.max_tokens_per_rank) + source_token) * Int32(self.topk) + source_slot token_in_expert = destination_position - pool_begin sf_position = sf_begin + token_in_expert diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/token_protocol.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/token_protocol.py index 289c5fe47298..e7f5e1252bfe 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/token_protocol.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/token_protocol.py @@ -11,19 +11,37 @@ @dataclasses.dataclass(frozen=True) class TokenSrcMetadata: - """One i64 routing record: rank:u16, topk:u16, token:u32.""" + """One i64 routing record: ``token:u16 | topk:u8 | rank:u8 | dst_order_id:u8 | (dst_token_cnt-1):u8 | reserved:u16``. + . ``dst_token_cnt`` = the token's routes landing on the destination rank, ``dst_order_id`` = this + route's position among them in ascending expert order (``order == cnt - 1`` is the last copy). Producers that do + not track copies use order 0, count 1.""" src_rank: Int32 src_token: Int32 src_topk: Int32 + dst_order_id: Int32 = Int32(0) + dst_token_cnt: Int32 = Int32(1) nbytes: ClassVar[int] = 8 + FIELD_BITS: ClassVar[int] = 8 + TOKEN_BITS: ClassVar[int] = 16 + MAX_RANKS: ClassVar[int] = 1 << 8 + MAX_TOPK: ClassVar[int] = 1 << 8 + MAX_TOKENS: ClassVar[int] = 1 << 16 def pack(self) -> Int64: - high = (Int64(self.src_rank) << Int64(16)) | Int64(self.src_topk) - return (high << Int64(32)) | ( - Int64(self.src_token) & Int64(0xFFFFFFFF) + high = Int64(self.dst_order_id) | ((Int64(self.dst_token_cnt) - Int64(1)) << Int64(8)) + low = ( + (Int64(self.src_rank) << Int64(24)) + | (Int64(self.src_topk) << Int64(16)) + | (Int64(self.src_token) & Int64(0xFFFF)) ) + return (high << Int64(32)) | low + + @property + def is_last_copy(self) -> Int32: + """1 when this row is the destination rank's last copy of the source token.""" + return Int32(self.dst_order_id == self.dst_token_cnt - Int32(1)) @staticmethod def _pointer(address: Union[cute.Pointer, Int64]) -> cute.Pointer: @@ -40,20 +58,24 @@ def _pointer(address: Union[cute.Pointer, Int64]) -> cute.Pointer: def store(self, address: Union[cute.Pointer, Int64]) -> None: cute.arch.store(self._pointer(address), self.pack(), scope="gpu") + @classmethod + def unpack(cls, packed: Int64) -> "TokenSrcMetadata": + """Inverse of ``pack`` for a record already in a register; every reader must decode through here.""" + high = packed >> Int64(32) + return cls( + src_rank=Int32((packed >> Int64(24)) & Int64(0xFF)), + src_token=Int32(packed & Int64(0xFFFF)), + src_topk=Int32((packed >> Int64(16)) & Int64(0xFF)), + dst_order_id=Int32(high & Int64(0xFF)), + dst_token_cnt=Int32((high >> Int64(8)) & Int64(0xFF)) + Int32(1), + ) + @classmethod def load( cls, address: Union[cute.Pointer, Int64], ) -> "TokenSrcMetadata": - packed = Int64( - cute.arch.load(cls._pointer(address), Int64, scope="gpu") - ) - high = packed >> Int64(32) - return cls( - src_rank=Int32((high >> Int64(16)) & Int64(0xFFFF)), - src_token=Int32(packed & Int64(0xFFFFFFFF)), - src_topk=Int32(high & Int64(0xFFFF)), - ) + return cls.unpack(Int64(cute.arch.load(cls._pointer(address), Int64, scope="gpu"))) __all__ = ["TokenSrcMetadata"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/ptx_helpers.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/ptx_helpers.py index 85d392151038..53192bb19bde 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/ptx_helpers.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/ptx_helpers.py @@ -495,6 +495,80 @@ def red_add_relaxed_sys_v2_bf16x2( ) +@dsl_user_op +def f32x2_pack(lo, hi, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None) -> Int64: + """Pack two f32 into one b64 f32x2 operand (low word ``lo``, high word ``hi``).""" + packed = llvm.inline_asm( + T.i64(), + [Float32(lo).ir_value(loc=loc, ip=ip), Float32(hi).ir_value(loc=loc, ip=ip)], + "mov.b64 $0, {$1, $2};", + "=l,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return Int64(packed) + + +@dsl_user_op +def fma_rn_f32x2_bf16x2( + word, scale2, acc2, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None +) -> Int64: + """``acc2 + word * scale2`` on a packed bf16x2 / f32x2 pair (``fma.rn.f32x2.bf16x2.f32x2.f32x2``, one rounding + per lane; scale 1.0 gives an exact ``acc + bf16``).""" + result = llvm.inline_asm( + T.i64(), + [ + Int32(word).ir_value(loc=loc, ip=ip), + Int64(scale2).ir_value(loc=loc, ip=ip), + Int64(acc2).ir_value(loc=loc, ip=ip), + ], + "fma.rn.f32x2.bf16x2.f32x2.f32x2 $0, $1, $2, $3;", + "=l,r,l,l", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return Int64(result) + + +@dsl_user_op +def cvt_rn_bf16x2_f32x2(acc2, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None) -> Int32: + """Round a packed f32x2 (low word -> low bf16 half) to one bf16x2 word with RN rounding (one F2FP).""" + packed = llvm.inline_asm( + T.i32(), + [Int64(acc2).ir_value(loc=loc, ip=ip)], + "{\n .reg .f32 lo, hi;\n mov.b64 {lo, hi}, $1;\n cvt.rn.bf16x2.f32 $0, hi, lo;\n}", + "=r,l", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return Int32(packed) + + +@dsl_user_op +def fence_proxy_async_global(*, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None) -> None: + """Order this thread's generic-proxy global accesses before its later async-proxy (TMA) accesses.""" + llvm.inline_asm( + None, + [], + "fence.proxy.async.global;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + @cute.jit def cvt_f32x4_to_f8x4_pack_i32(fp32x4: cute.Tensor, fp8_type, *, loc=None, ip=None) -> Int32: """Round four f32 lanes to the selected FP8 format and pack them into one i32.""" @@ -759,4 +833,9 @@ def stg_e8m0x8_from_f32( "stg_e8m0x8_from_f32", "stg_f32", "tma_load_1d", + "fence_acq_rel_gpu", + "fence_proxy_async_global", + "f32x2_pack", + "fma_rn_f32x2_bf16x2", + "cvt_rn_bf16x2_f32x2", ] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py index 2d503d43c52f..fe912753b951 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py @@ -33,10 +33,12 @@ from .....helpers.cute_py_helpers import tcgen05_block_scaled_acc_dtype from .....helpers.dsl_helpers import mark_alignment from .....helpers.flag_batch import make_flag_batch_tracker +from .....helpers.iket_compat import iket from .....helpers.ptx_helpers import ( cp_async_bulk_s2g, cp_reduce_async_bulk_add_bf16_s2g, cvt_f32_to_fp8_to_f32, + fence_proxy_async_global, movmatrix_b16, red_add_relaxed_sys_v2_bf16x2, ) @@ -920,6 +922,10 @@ def impl_desc_require(cls) -> dict[str, object]: "token_back_ready_granularity": OptionalRequirement(str), } + def _make_fc2_epilogue(self, *args) -> "SwapABFc2Epilogue": + """Device-side FC2 epilogue factory; variants substitute their own ``SwapABFc2Epilogue`` subclass.""" + return SwapABFc2Epilogue(self, *args) + @staticmethod def _resolve_situ_betas(problem_desc: ProblemDesc) -> Tuple[Optional[float], Optional[float]]: """Resolve the SiTU (Kimi K3) betas from the ProblemDesc; ``(None, None)`` means SwiGLU. @@ -1286,8 +1292,7 @@ def run( fc1_done_counter, optional_epi_args, ) - fc2_epi = SwapABFc2Epilogue( - self, + fc2_epi = self._make_fc2_epilogue( tidx, fc2_smem_tensor, tma_atom_fc2_output, @@ -2112,12 +2117,17 @@ def run_subtile( alpha_val: Optional[cutlass.Float32], ): process_pipeline = self.process_pipeline + iket.range_push("fc2.tmem_load") loaded = process_pipeline.tmem_acc_load(tmem_subtile_tensor=tmem_subtile_tensor, epi=self) + iket.range_pop() + iket.range_push("fc2.f2fp_reorder") casted = process_pipeline.f2fp(*loaded, alpha_val=alpha_val) # reorder returns a bare RMEM fragment in the store's expected pre-store # distribution; reorder + store are paired 1:1 inside the pipeline. pre_store = process_pipeline.post_f2fp_reorder(casted=casted, tmem_subtile_view=tmem_subtile_tensor) + iket.range_pop() + iket.range_push("fc2.store") process_pipeline.store_function( epi=self, subtile=pre_store, @@ -2126,6 +2136,7 @@ def run_subtile( subtile_idx=subtile_idx, fc2_output_router=fc2_output_router, ) + iket.range_pop() @dataclasses.dataclass(frozen=True) @@ -2603,6 +2614,9 @@ def fc2_stg_store_function( src_i = cute.make_tensor(copy_src[None, local_iter].iterator, single_copy_layout) dst_i = cute.make_tensor(dst_ptr, single_copy_layout) cute.copy(copy_atom_vec, cute.recast_tensor(src_i, cutlass.Int32), cute.recast_tensor(dst_i, cutlass.Int32)) + if cutlass.const_expr(epi.token_back_push_data): + # Generic-proxy pool rows are read back by the combine warps' TMA: fence before the fc2_done release. + fence_proxy_async_global() @cute.jit diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue_dedup.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue_dedup.py new file mode 100644 index 000000000000..5364aa737f04 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue_dedup.py @@ -0,0 +1,349 @@ +"""Dedup FC2 epilogue: SMEM-staged subtiles pushed with one ``cp.async.bulk`` per token. ``dst_token_cnt == 1`` +rows go straight to the source rank's ``pre_reduced[src_token][my_rank]``, others into the local order buffer +``[T, R, chunks_per_row, K, chunk]`` for the combine warps. Top-k scores are applied at FC1 only.""" + +from __future__ import annotations + +import dataclasses +from typing import Optional, Tuple + +import cutlass +import cutlass.cute as cute +from cutlass.cute.typing import AddressSpace +from cutlass.cutlass_dsl import Int32, Int64 + +from .....api import ImplDesc, ProblemDesc +from .....communication.nvlink_domain.symmetric_buffer import SymmetricBufferDevice +from .....communication.token_protocol import TokenSrcMetadata +from .....helpers.iket_compat import iket +from .....helpers.ptx_helpers import ( + cp_async_bulk_s2g, + cp_reduce_async_bulk_add_bf16_s2g, +) +from ....function_mapping import FunctionMapping +from ....schedulers import SwapAbFc12WorkTileInfo +from .block_scaled_swap_ab_fc12_epilogue import ( + GatedActEpilogueArgs, + SwapABFc2Epilogue, + SwapABGatedActEpilogue, + make_fc2_ublk_process_pipeline, +) + + +@dataclasses.dataclass(frozen=True) +class GatedActEpilogueArgsDedup(GatedActEpilogueArgs): + """Base epilogue args plus the pre-reduced (T, R, H) buffer and this rank's index (the kernel's host argument, + which selects the ``pre_reduced[src_token][local_rank]`` slot of single-copy rows).""" + + fc2_remote_output: Optional[cute.Tensor] = None + local_rank: Optional[Int32] = None + + +@dataclasses.dataclass(frozen=True) +class Fc2DedupBulkRouter: + """Per work tile: the bulk-copy destination of every token this thread owns. Order-buffer slot = + ``(((src_token * R + src_rank) * chunks_per_row + hidden_chunk) * K + dst_order_id) * chunk``; single-copy + tokens map to ``pre_reduced[src_token][slot][hidden_base]`` (slot = my rank, 0 under in-kernel REDG).""" + + metadata: cute.Tensor # domain-offset to the tile's first pool row + order_base: Int64 + remote_output: cute.Tensor # local pre-reduced buffer (T, R, H); (T, 1, H) view of the output under REDG + peer_rank_ptr_mapper: SymmetricBufferDevice + local_rank: Int32 + hidden_base: Int32 # first hidden element of this CTA tile + valid_tokens_this_cta_tile: Int32 + valid_hidden_this_cta_tile: Int32 + world_size: int + topk: int + chunk_bytes: int + chunks_per_row: int + element_bytes: int + reduce_topk_in_kernel: bool + mapping: FunctionMapping # per-issue (epi_tid, iter_idx) -> token_in_cta_tile + epi_tid: Int32 + dst_ptrs: Optional[cute.Tensor] = None # i64 x copy_iters + valid: Optional[cute.Tensor] = None # i32 x copy_iters: 0 skip, 1 bulk copy, 2 bulk reduce-add + copy_iters: int = 0 + + @cute.jit + def prefetch(self) -> "Fc2DedupBulkRouter": + copy_iters: cutlass.Constexpr[int] = self.mapping.domain.axis_size("iter_idx") + valid = cute.make_rmem_tensor((copy_iters,), Int32) + dst_ptrs = cute.make_rmem_tensor((copy_iters,), Int64) + hidden_base_bytes = Int64(self.hidden_base) * Int64(self.element_bytes) + chunk_index = hidden_base_bytes // Int64(self.chunk_bytes) + byte_in_chunk = hidden_base_bytes - chunk_index * Int64(self.chunk_bytes) + remote_slot = Int32(0) if cutlass.const_expr(self.reduce_topk_in_kernel) else self.local_rank + for iter_idx in cutlass.range_constexpr(copy_iters): + token_in_tile = Int32(self.mapping.evaluate(epi_tid=self.epi_tid, iter_idx=iter_idx)["token_in_cta_tile"]) + valid[iter_idx] = Int32(0) + dst_ptrs[iter_idx] = Int64(0) + if (token_in_tile < self.valid_tokens_this_cta_tile) and (self.valid_hidden_this_cta_tile > Int32(0)): + md = TokenSrcMetadata.load( + self.metadata.iterator.toint() + Int64(token_in_tile) * Int64(TokenSrcMetadata.nbytes) + ) + if md.dst_token_cnt == Int32(1): + # Single copy on this rank: push straight to the source rank's pre-reduced row. + local_row = cute.domain_offset( + (Int64(md.src_token), remote_slot, self.hidden_base), self.remote_output + ).iterator.toint() + dst_ptrs[iter_idx] = self.peer_rank_ptr_mapper.map(local_row, md.src_rank) + valid[iter_idx] = Int32(2) if cutlass.const_expr(self.reduce_topk_in_kernel) else Int32(1) + else: + block_index = Int64(md.src_token) * Int64(self.world_size) + Int64(md.src_rank) + slot_index = (block_index * Int64(self.chunks_per_row) + chunk_index) * Int64(self.topk) + Int64( + md.dst_order_id + ) + dst_ptrs[iter_idx] = self.order_base + slot_index * Int64(self.chunk_bytes) + byte_in_chunk + valid[iter_idx] = Int32(1) + return dataclasses.replace(self, dst_ptrs=dst_ptrs, valid=valid, copy_iters=copy_iters) + + @cute.jit + def get_data_dst(self, iter_idx: int) -> Tuple[cute.Pointer, Int32]: + ptr = cute.make_ptr(cutlass.BFloat16, self.dst_ptrs[iter_idx], AddressSpace.gmem, assumed_align=16) + return ptr, self.valid[iter_idx] + + +@cute.jit +def fc2_dedup_bulk_store_function( + *, + epi, + subtile: cute.Tensor, # (64,) bf16: one element per token of the subtile (hidden = warp*32 + lane) + epilogue_iter_idx: Int32, + subtile_idx: Int32, + fc2_output_router: Fc2DedupBulkRouter, + **_, +): + smem_tensor = epi.smem_tensor + if cutlass.const_expr(smem_tensor is None): + raise ValueError("The dedup bulk store requires epi.smem_tensor.") + data_subtile = subtile + if cutlass.const_expr(data_subtile.element_type is not cutlass.BFloat16): + raise ValueError("The dedup FC2 store expects a BF16 payload.") + if cutlass.const_expr(cute.size(data_subtile) != epi._EpilogueTokenTileSize): + raise ValueError("BF16 bulk staging requires one register per token in the 64-token subtile.") + warp_idx = epi.warp_idx + lane_idx = epi.lane_idx + smem_read_write_bar = SwapABGatedActEpilogue.epilogue_sync_barrier() + stage_idx = epilogue_iter_idx % Int32(epi.fc2_tma_stages) + smem_stage = smem_tensor[None, None, stage_idx] + + warp_hidden_base = Int32(warp_idx * 32) + for token_idx in cutlass.range_constexpr(epi._EpilogueTokenTileSize): + smem_stage[token_idx, warp_hidden_base + lane_idx] = data_subtile[token_idx] + + cute.arch.fence_proxy("async.shared", space="cta") + smem_read_write_bar.arrive_and_wait() + + copy_elems = Int32(epi._EpilogueFc2HiddenTileSize) + if cutlass.const_expr(epi.fc2_hidden_needs_predicate): + copy_elems = fc2_output_router.valid_hidden_this_cta_tile + copy_bytes = copy_elems * Int32(epi.combine_format.act_dtype.width // 8) + scratch_row = warp_idx * Int32(16) + lane_idx % Int32(16) + copy_iters: cutlass.Constexpr[int] = fc2_output_router.mapping.domain.axis_size("iter_idx") + for ublk_iter_idx in cutlass.range_constexpr(copy_iters): + owned_subtile = Int32(ublk_iter_idx * 2) + lane_idx // Int32(16) + if owned_subtile == subtile_idx: + dst_ptr, kind = fc2_output_router.get_data_dst(ublk_iter_idx) + if kind != Int32(0): + src_row = cute.slice_(smem_stage, (scratch_row, None)) + if cutlass.const_expr(fc2_output_router.reduce_topk_in_kernel): + if kind == Int32(2): + cp_reduce_async_bulk_add_bf16_s2g(dst_ptr, src_row.iterator, copy_bytes) + else: + cp_async_bulk_s2g(dst_ptr, src_row.iterator, copy_bytes) + else: + cp_async_bulk_s2g(dst_ptr, src_row.iterator, copy_bytes) + + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(epi.fc2_tma_stages - 1, read=True) + smem_read_write_bar.arrive_and_wait() + + +class SwapABFc2EpilogueDedup(SwapABFc2Epilogue): + """FC2 epilogue with SMEM-staged, metadata-routed bulk stores (order buffer or direct peer push).""" + + def __init__( + self, + base: "SwapABGatedActEpilogueDedup", + tidx: Int32, + smem_tensor: Optional[cute.Tensor], + tma_atom_fc2_output: Optional[cute.CopyAtom], + fc2_tma_output: Optional[cute.Tensor], + fc2_output: cute.Tensor, # order buffer base as a (T*R*K, 1, hidden) tensor + token_src_metadata: Optional[cute.Tensor], + fc2_done_counter: Optional[cute.Tensor], + fc2_output_sf: Optional[cute.Tensor], + peer_rank_ptr_mapper, + optional_epi_args: GatedActEpilogueArgs, + ): + # Not super().__init__(): the base picks its store pipeline and freezes; dedup installs its own. + self.base = base + self.tidx = tidx % (base._EpilogueWarpCnt * 32) + self.warp_idx = self.tidx // 32 + self.lane_idx = self.tidx % 32 + self.tma_atom_fc2_output = None + self.fc2_tma_output = None + self.fc2_output = fc2_output + self.token_src_metadata = token_src_metadata + self.fc2_done_counter = fc2_done_counter + self.fc2_output_sf = None + self.peer_rank_ptr_mapper = peer_rank_ptr_mapper + self.optional_epi_args = optional_epi_args + if smem_tensor is None: + raise ValueError("The dedup FC2 epilogue needs the FC2 staging SMEM (fc2_use_bulk).") + self.smem_tensor = smem_tensor + self.process_pipeline = dataclasses.replace( + make_fc2_ublk_process_pipeline( + combine_format=base.combine_format, + cta_token_tile_size=base.cta_tile_n, + cta_hidden_tile_size=base.cta_tile_m, + ), + store_function=fc2_dedup_bulk_store_function, + ) + self._freeze() + + @cute.jit + def _make_bulk_router(self, work_tile_info: SwapAbFc12WorkTileInfo) -> Fc2DedupBulkRouter: + task_tile_data_row_start = ( + work_tile_info.cumulative_data_physical_row + work_tile_info.tile_n_idx * Int32(self.cta_tile_n) + ) + hidden_base_this_cta_tile = work_tile_info.tile_m_idx * Int32(self.cta_tile_m) + valid_hidden = Int32(self.fc2_output.shape[2]) - hidden_base_this_cta_tile + if valid_hidden < 0: + valid_hidden = 0 + if valid_hidden > self._EpilogueFc2HiddenTileSize: + valid_hidden = self._EpilogueFc2HiddenTileSize + remote_output = self.optional_epi_args.fc2_remote_output + if cutlass.const_expr(remote_output is None): + raise ValueError("The dedup FC2 epilogue needs fc2_remote_output (the pre-reduced (T, R, H) buffer).") + # The pre_reduced slot of single-copy rows is the kernel's local_rank argument (not inferred from the mapper). + local_rank = self.optional_epi_args.local_rank + if cutlass.const_expr(local_rank is None): + raise ValueError("The dedup FC2 epilogue needs local_rank (the kernel's host-provided rank argument).") + return Fc2DedupBulkRouter( + metadata=cute.domain_offset((task_tile_data_row_start,), self.token_src_metadata), + order_base=self.fc2_output.iterator.toint(), + remote_output=remote_output, + peer_rank_ptr_mapper=self.peer_rank_ptr_mapper, + local_rank=Int32(local_rank), + hidden_base=hidden_base_this_cta_tile, + valid_tokens_this_cta_tile=work_tile_info.valid_tokens_in_cta_tile, + valid_hidden_this_cta_tile=valid_hidden, + world_size=self.dedup_world_size, + topk=self.dedup_topk, + chunk_bytes=self.dedup_chunk_bytes, + chunks_per_row=self.dedup_chunks_per_row, + element_bytes=self.combine_format.act_dtype.width // 8, + reduce_topk_in_kernel=self.reduce_topk_in_kernel, + mapping=self.process_pipeline.store_out_mapping, + epi_tid=self.tidx, + ) + + @cute.jit + def __call__( + self, work_tile_info: SwapAbFc12WorkTileInfo, tmem_acc_tensor: cute.Tensor, acc_pipeline, acc_consumer_state + ): + if cutlass.const_expr(self.optional_epi_args.fc2_alpha is not None): + alpha_val = self.optional_epi_args.fc2_alpha[work_tile_info.expert_idx] + else: + alpha_val = None + router = self._make_bulk_router(work_tile_info) + acc_ready = False + if not work_tile_info.peek_ready: + acc_ready = True + acc_pipeline.consumer_wait(acc_consumer_state) + tmem_acc_tensor_tiled_by_epi_tile = cute.flat_divide( + tmem_acc_tensor, (self._EpilogueFc2HiddenTileSize, self._EpilogueTokenTileSize) + )[None, None, 0, None] + acc_pipeline.consumer_wait(acc_consumer_state, acc_ready) + # Read the metadata only after the accumulator wait: under PDL the router may still be pushing records. + router = router.prefetch() + valid_tokens = work_tile_info.valid_tokens_in_cta_tile + for subtile_idx in cutlass.range(self.subtile_cnt, unroll=1): + if subtile_idx * Int32(self._EpilogueTokenTileSize) < valid_tokens: + self.run_subtile( + work_tile_info=work_tile_info, + epilogue_iter_idx=subtile_idx, + subtile_idx=subtile_idx, + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx], + fc2_output_router=router, + alpha_val=alpha_val, + ) + cute.arch.fence_view_async_tmem_load() + acc_pipeline.consumer_release(acc_consumer_state) + + +class SwapABGatedActEpilogueDedup(SwapABGatedActEpilogue): + """``SwapABGatedActEpilogue`` whose FC2 half is ``SwapABFc2EpilogueDedup`` (bulk stores routed by metadata).""" + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return {**super().problem_desc_require(), "apply_topk_at_fc1": bool} + + @classmethod + def impl_desc_require(cls) -> dict[str, object]: + return { + **super().impl_desc_require(), + "dedup_world_size": int, + "dedup_topk": int, + # Order buffer layout [T, R, chunks_per_row, K, chunk_bytes] (from TokenCommDedup). + "dedup_chunk_bytes": int, + "dedup_chunks_per_row": int, + } + + def _make_fc2_epilogue(self, *args) -> SwapABFc2EpilogueDedup: + return SwapABFc2EpilogueDedup(self, *args) + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + super().__init__(problem_desc, impl_desc) + self.dedup_world_size = int(impl_desc["dedup_world_size"]) + self.dedup_topk = int(impl_desc["dedup_topk"]) + self.dedup_chunk_bytes = int(impl_desc["dedup_chunk_bytes"]) + self.dedup_chunks_per_row = int(impl_desc["dedup_chunks_per_row"]) + self.apply_topk_at_fc1 = bool(problem_desc["apply_topk_at_fc1"]) + if not self.apply_topk_at_fc1: + raise ValueError("Combine dedup applies top-k scores at FC1 only (pre_fc1_quant); post_fc2 is not supported.") + if not self.communication_enabled: + raise ValueError("Combine dedup requires token communication.") + if not self.token_back_push_data: + raise ValueError( + "Combine dedup reduces and pushes the order-buffer blocks from the combine warps; epi_warps is not supported." + ) + if not self.fc2_use_bulk: + raise ValueError("Combine dedup stages FC2 subtiles in SMEM for per-token bulk stores; fc2_use_bulk must be on.") + if self.combine_format.is_quantized or self.combine_format.act_dtype is not cutlass.BFloat16: + raise ValueError("Combine dedup supports the BF16 combine format only.") + if self.reduce_topk_in_epilogue: + raise ValueError("Combine dedup never reduces in the epilogue.") + if self.dedup_world_size <= 0 or self.dedup_topk <= 0 or self.dedup_chunks_per_row <= 0: + raise ValueError("Combine dedup needs positive dedup_world_size, dedup_topk and dedup_chunks_per_row.") + tile_bytes = self._EpilogueFc2HiddenTileSize * self.combine_format.act_dtype.width // 8 + if self.dedup_chunk_bytes <= 0 or self.dedup_chunk_bytes % tile_bytes != 0: + raise ValueError( + f"Combine dedup needs dedup_chunk_bytes as a multiple of the {tile_bytes} B hidden tile, got " + f"{self.dedup_chunk_bytes}." + ) + # The base resolved fc2_use_bulk + local output as the TMA tile store; dedup stores per token with UBLK. + self.fc2_use_tma = False + self.fc2_use_ublk = True + single_stage_region = self._make_fc2_single_stage_region() + if single_stage_region.nbytes % 16 != 0: + raise ValueError("Each FC2 UBLK staging stage must occupy a multiple of 16 bytes.") + requested_fc2_tma_stages = impl_desc.get("fc2_tma_stages") + if requested_fc2_tma_stages is not None: + self.fc2_tma_stages = int(requested_fc2_tma_stages) + else: + bf16_stage_bytes = self._EpilogueTokenTileSize * self._EpilogueFc2HiddenTileSize * cutlass.BFloat16.width // 8 + available_staging_bytes = max(self.fc1_staging_bytes, min(2, self.subtile_cnt) * bf16_stage_bytes) + self.fc2_tma_stages = min(self.subtile_cnt, available_staging_bytes // single_stage_region.nbytes) + self.fc2_staging_spec = self._make_fc2_staging_region(self.fc2_tma_stages) + + +__all__ = [ + "Fc2DedupBulkRouter", + "GatedActEpilogueArgsDedup", + "SwapABFc2EpilogueDedup", + "SwapABGatedActEpilogueDedup", + "fc2_dedup_bulk_store_function", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_mainloop.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_mainloop.py index bcd101a404fe..b8fd0871d71f 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_mainloop.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_mainloop.py @@ -241,6 +241,8 @@ def _validate_configuration(self) -> None: def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: """Register the selected Rubin mainloop SMEM plan.""" smem_workspace.register_mbarrier(self.ab_pipeline_mbarriers_region, self.num_ab_pipeline_stages * 2) + # Register the 4 B holding word before the 128 B-aligned stage tensors so it lands in the alignment gap. + smem_workspace.register_tensor(self.tmem_holding_buffer_region, cutlass.Int32, (1,), byte_alignment=8) for region_name, region in ( (self.a_smem_tensor_region, self.a_smem_region), (self.b_smem_tensor_region, self.b_smem_region), @@ -257,7 +259,6 @@ def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: ) smem_workspace.register_mbarrier(self.acc_pipeline_mbarriers_region, self.num_accumulator_pipeline_stages * 2) smem_workspace.register_mbarrier(self.tmem_deallocation_mbarrier_region, 1) - smem_workspace.register_tensor(self.tmem_holding_buffer_region, cutlass.Int32, (1,), byte_alignment=8) def _select_ab_stage_plan(self) -> None: stage_one_regions = make_smem_layouts( @@ -287,7 +288,9 @@ def packed_bytes(region, dtype) -> int: ) mbarrier_bytes = int(cutlass.Int64.width) // 8 ab_stage_cost_bytes = self.ab_stage_payload_bytes + 2 * mbarrier_bytes - plan_tail_bytes = (2 * self.num_accumulator_pipeline_stages + 2) * mbarrier_bytes + # Accumulator + TMEM-deallocation mbarriers plus the 8 B-aligned TMEM holding word. + tmem_holding_buffer_bytes = 8 + plan_tail_bytes = (2 * self.num_accumulator_pipeline_stages + 2) * mbarrier_bytes + tmem_holding_buffer_bytes self.num_ab_pipeline_stages = (self.mainloop_smem_budget_bytes - plan_tail_bytes) // ab_stage_cost_bytes if self.num_ab_pipeline_stages < 1: raise ValueError( diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py index 3c5bb77f999c..b8d4910a0cdc 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py @@ -4,6 +4,8 @@ from typing import ClassVar, Optional, Tuple +import os + import cuda.bindings.driver as cuda import cutlass @@ -21,6 +23,8 @@ derive_helper_weight_ready_topology, make_helper_weight_ready_generation_gate, ) +from .....communication.nvlink_domain.token_comm_dedup import TokenCommDedup +from .....communication.token_protocol import TokenSrcMetadata from .....helpers.cute_py_helpers import Tcgen05MmaInstruction, make_tcgen05_tmem_plan, tcgen05_block_scaled_acc_dtype from .....helpers.device_workspace import DeviceWorkspace from .....helpers.dlb_expert_count import resolve_dlb_expert_counts @@ -40,8 +44,10 @@ from ....schedulers.non_clc_mixed_cga import NonClcMixedCgaConfig from ...custom_mix_cga_helpers import TmaAtomOrPair, pipeline_init_arrive_mixed_cga, pipeline_init_wait_mixed_cga from .block_scaled_swap_ab_fc12_epilogue import GatedActEpilogueArgs, SwapABGatedActEpilogue +from .block_scaled_swap_ab_fc12_epilogue_dedup import GatedActEpilogueArgsDedup, SwapABGatedActEpilogueDedup from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension from .block_scaled_swap_ab_fc12_mainloop import BlockScaledSwapAbFc12Mainloop +from .topk_reduce_dedup import TopkReduceDedup _aot_symbol_prefix = "rubin_mega_moe_aot" @@ -117,6 +123,7 @@ def impl_desc_require(cls) -> dict[str, type]: "reduce_topk_in_kernel": bool, "helper_expert_count": OptionalRequirement(int), "total_helper_slots": OptionalRequirement(int), + "combine_dedup": OptionalRequirement(bool), } def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: @@ -192,6 +199,10 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.helper_expert_count = impl_desc.get("helper_expert_count", 0) if type(self.helper_expert_count) is not int or self.helper_expert_count < 0: raise ValueError("helper_expert_count must be a non-negative exact int") + self.combine_dedup = bool(impl_desc.get("combine_dedup", False)) + # Diagnostic only (MEGA_DIAG_SKIP_TOKEN_BACK=1): combine warps do nothing, so the timed kernel is + # MMA + epilogue + dispatch alone. Output is wrong; run with MEGA_TIMING_ONLY=1. + self._diag_skip_token_back = os.environ.get("MEGA_DIAG_SKIP_TOKEN_BACK", "0") == "1" self.occupancy = 1 self.architecture = "sm_107" self.local_expert_count = self.expert_count // self.world_size @@ -206,8 +217,9 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: # static M is required only when an S>0 helper boundary must be # compiled into the scheduler gate. self.home_expert_count = self.local_expert_count - self.threads_per_cta = 16 * 32 if self.token_back_mode == "standalone_warps" else 12 * 32 # nosec B105 - self.other_warp_register_count = 64 if self.token_back_mode == "standalone_warps" else 72 # nosec B105 + n_warps = 16 if self.token_back_mode in ("standalone_warps", "hybrid_warps") else 12 + self.threads_per_cta = n_warps * 32 + self.other_warp_register_count = 64 if self.token_back_mode in ("standalone_warps", "hybrid_warps") else 72 self.mixed_cga_config = NonClcMixedCgaConfig( preferred_cluster_shape=self.cluster_shape_mn, @@ -276,7 +288,11 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: "router_smem_limit_bytes": utils.get_smem_capacity_in_bytes(self.architecture), } ) - self.token_comm = TokenCommNonDeterministic(problem_desc, token_comm_impl_desc) + # One top-level assignment on purpose: tests/test_dlb_contracts.py replays the constructor's + # descriptor-forwarding statements and only sees top-level assignments to self.token_comm. + self.token_comm = (TokenCommDedup if self.combine_dedup else TokenCommNonDeterministic)( + problem_desc, token_comm_impl_desc + ) self.max_tokens = self.token_comm.worst_case_token_count fc12_problem_desc = ProblemDesc({**problem_desc, "expert_count": self.local_expert_count}) @@ -298,13 +314,28 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: "fc2_epi_flag_batch": self.epi_flag_batches[1], "compute_cluster_count": self.compute_cluster_count, "token_back_ready_granularity": self.token_back_ready_granularity, + # Combine dedup: the dedup FC2 epilogue scatters rows into the local order buffer + # [T, R, chunks_per_row, K, chunk] whose geometry TokenCommDedup owns. + **( + { + "dedup_world_size": self.world_size, + "dedup_topk": self.topk, + "dedup_chunk_bytes": self.token_comm.order_chunk_bytes, + "dedup_chunks_per_row": self.token_comm.order_chunks_per_row, + } + if self.combine_dedup + else {} + ), } ) scheduler_type = ( BlackwellFusedFc12Scheduler if self.schedule_mode == "grouped" else PhaseInterleavedFc12Scheduler ) self.scheduler = scheduler_type(fc12_problem_desc, resolved_impl_desc) - self.epilogue = SwapABGatedActEpilogue(fc12_problem_desc, resolved_impl_desc) + if self.combine_dedup: + self.epilogue = SwapABGatedActEpilogueDedup(fc12_problem_desc, resolved_impl_desc) + else: + self.epilogue = SwapABGatedActEpilogue(fc12_problem_desc, resolved_impl_desc) if self.token_back_ready_granularity == "token_tile": # nosec B105 self._validate_token_tile_ready_contract() if self.epilogue.fc1_output_dtype is not self.b_dtype: @@ -340,9 +371,12 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: if self._comm_smem_workspace is not None: self._comm_smem_workspace.assert_fits_arena(self.smem_arena_bytes) self._smem_workspace.assert_fits_arena(self.smem_arena_bytes) - self._topk_reduce = ( - None if self.reduce_topk_in_kernel else TopkReduce(self.hidden_size, self.topk, self.combine_format) - ) + if self.reduce_topk_in_kernel: + self._topk_reduce = None + elif self.combine_dedup: + self._topk_reduce = TopkReduceDedup(self.hidden_size, self.world_size, self.combine_format) + else: + self._topk_reduce = TopkReduce(self.hidden_size, self.topk, self.combine_format) def _validate_cta_partition(self) -> None: for name, value in (("launch_cluster_count", self.launch_cluster_count), ("comm_clusters", self.comm_clusters)): @@ -475,8 +509,37 @@ def _validate_geometry(self) -> None: raise ValueError("topk must be positive and no greater than expert_count.") if self.topk_index_dtype not in (cutlass.Int32, cutlass.Int64): raise ValueError(f"topk_index_dtype must be Int32 or Int64, got {self.topk_index_dtype}.") + if self.topk > TokenSrcMetadata.MAX_TOPK or self.world_size > TokenSrcMetadata.MAX_RANKS: + raise ValueError( + f"TokenSrcMetadata keeps topk and rank in {TokenSrcMetadata.FIELD_BITS} bits each; " + f"topk must be <= {TokenSrcMetadata.MAX_TOPK} and world_size <= {TokenSrcMetadata.MAX_RANKS}." + ) + if self.max_tokens_per_rank > TokenSrcMetadata.MAX_TOKENS: + raise ValueError( + f"TokenSrcMetadata keeps the token index in {TokenSrcMetadata.TOKEN_BITS} bits; " + f"max_tokens_per_rank must be <= {TokenSrcMetadata.MAX_TOKENS}." + ) if self.reduce_topk_in_kernel and not self.apply_topk_at_fc1: raise ValueError("In-kernel top-k reduction requires apply_topk_at_fc1=True.") + if self.combine_dedup: + if not self.apply_topk_at_fc1: + raise ValueError("combine_dedup applies top-k scores at FC1 only (pre_fc1_quant).") + if self.combine_format.is_quantized or self.combine_format.act_dtype is not cutlass.BFloat16: + raise ValueError("combine_dedup supports the BF16 combine format only.") + if self.world_size > 32: + raise ValueError("combine_dedup keeps the destination-rank mask in an Int32; world_size must be <= 32.") + if self.token_back_mode == "epi_warps": # nosec B105 + raise ValueError( + "combine_dedup pushes the reduced rank-major rows from the combine warps; use " + "reuse_dispatch_warps or standalone_warps." + ) + if not self.fc2_use_bulk: + raise ValueError("combine_dedup stages FC2 subtiles in SMEM for per-token bulk stores; fc2_use_bulk must be on.") + if self.comm_clusters > 0: + raise ValueError("combine_dedup does not support the communication CTA split (comm_clusters must be 0).") + if self.token_back_ready_granularity != "expert": # nosec B105 + # The dedup combine warps poll one fc2_done counter per local expert (expert-prefix wait). + raise ValueError("combine_dedup signals FC2 completion per expert; token_back_ready_granularity must be 'expert'.") if self.intermediate_gateup_size % 2 != 0: raise ValueError("The SwiGLU intermediate dimension must be even.") intermediate_downproj = self.intermediate_gateup_size // 2 @@ -689,6 +752,8 @@ def dtype_name(dtype: type) -> str: f"{'inkernel_redg' if self.reduce_topk_in_kernel else 'separate_reduce'}" "_dlb_v1" + ("_fc2readytile_v1" if self.token_back_ready_granularity == "token_tile" else "") # nosec B105 + + f"_{'dedup' if self.combine_dedup else 'nodedup'}" + + ("_diagskiptb" if self._diag_skip_token_back else "") ) def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): @@ -1020,8 +1085,13 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: else: pre_reduced_activation = self.token_comm.pre_reduced_activation_tensor(self._device_workspace) pre_reduced_activation_sf = self.token_comm.pre_reduced_activation_sf_tensor(self._device_workspace) + dedup_valid_rank_mask = None + if cutlass.const_expr(self.combine_dedup): + dedup_valid_rank_mask = self.token_comm.valid_rank_mask_tensor(self._device_workspace) if cutlass.const_expr(self.token_comm.token_back_push_data): + # Local FC2 output the combine warps push from: the local-expert-major pool, or under dedup the + # (T*R*K, 1, H) order buffer the STG store scatters into by metadata. fc2_output = self.token_comm.fc2_activation_tensor(self._device_workspace) else: fc2_output = rewrite_tensor_shape( @@ -1126,10 +1196,14 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: use_pdl=True, ) if cutlass.const_expr(not self.reduce_topk_in_kernel): - reduce_scores = None if cutlass.const_expr(self.apply_topk_at_fc1) else topk_scores - self._topk_reduce( - pre_reduced_activation, pre_reduced_activation_sf, output_activation, reduce_scores, stream - ) + if cutlass.const_expr(self.combine_dedup): + # Weights are already folded (FC1 pre-multiply or FC2-epilogue multiply); sum the rank slots. + self._topk_reduce(pre_reduced_activation, dedup_valid_rank_mask, output_activation, stream) + else: + reduce_scores = None if cutlass.const_expr(self.apply_topk_at_fc1) else topk_scores + self._topk_reduce( + pre_reduced_activation, pre_reduced_activation_sf, output_activation, reduce_scores, stream + ) @cute.kernel @@ -1270,9 +1344,24 @@ def _kernel( fc2_spin_threshold=fc2_spin_threshold, fc1_ready_counter_pointer=self.token_comm.fc1_ready_counter_pointer(self._device_workspace), ) - optional_epilogue_args = GatedActEpilogueArgs( - fc1_alpha=fc1_alpha, fc2_alpha=fc2_alpha, fc1_norm_const=fc1_norm_const, topk_scores=pool_topk_scores - ) + # FC1 pre-multiplies under apply_topk_at_fc1 (the only position combine_dedup supports). + fc1_pool_topk_scores = pool_topk_scores if cutlass.const_expr(self.apply_topk_at_fc1) else None + if cutlass.const_expr(self.combine_dedup): + optional_epilogue_args = GatedActEpilogueArgsDedup( + fc1_alpha=fc1_alpha, + fc2_alpha=fc2_alpha, + fc1_norm_const=fc1_norm_const, + topk_scores=fc1_pool_topk_scores, + fc2_remote_output=pre_reduced_activation, + local_rank=local_rank, + ) + else: + optional_epilogue_args = GatedActEpilogueArgs( + fc1_alpha=fc1_alpha, + fc2_alpha=fc2_alpha, + fc1_norm_const=fc1_norm_const, + topk_scores=fc1_pool_topk_scores, + ) pipeline_init_arrive_mixed_cga( self.cluster_shape_mn, self.resolved_fallback_cluster_shape_mn, is_fallback_cluster, is_relaxed=True @@ -1280,7 +1369,9 @@ def _kernel( pipeline_init_wait_mixed_cga( self.cluster_shape_mn, self.resolved_fallback_cluster_shape_mn, is_fallback_cluster ) - finalize_barrier_thread_count = 12 * 32 if self.token_back_mode == "standalone_warps" else 8 * 32 # nosec B105 + finalize_barrier_thread_count = ( + 12 * 32 if self.token_back_mode in ("standalone_warps", "hybrid_warps") else 8 * 32 + ) finalize_barrier = pipeline.NamedBarrier(barrier_id=13, num_threads=finalize_barrier_thread_count) if warp_idx == self.scheduler_warp_id: @@ -1409,8 +1500,15 @@ def _kernel( iket.range_push("mega.token_in") self.token_comm.token_in(self._smem_workspace, smem_base) iket.range_pop() + # epi_warps pushes the data from the epilogue, but a quantized wire still needs its scale rows pushed + # token-contiguously (token_back_push_sf); the dispatch warps do that after token_in. if cutlass.const_expr( - self.token_comm.token_back_enabled and self.token_back_mode != "standalone_warps" # nosec B105 + self.token_comm.token_back_enabled + and ( + self.token_back_mode in ("reuse_dispatch_warps", "hybrid_warps") + or (self.token_back_mode == "epi_warps" and not self.token_comm.token_back_push_data) # nosec B105 + ) + and not self._diag_skip_token_back ): iket.range_push("mega.token_back") self.token_comm.token_back(self._smem_workspace, smem_base) @@ -1419,9 +1517,10 @@ def _kernel( iket.range_push("mega.tail_reset") self.token_comm.reset_tail() iket.range_pop() - elif cutlass.const_expr(self.token_back_mode == "standalone_warps"): # nosec B105 + elif cutlass.const_expr(self.token_back_mode in ("standalone_warps", "hybrid_warps")): iket.range_push("mega.token_back_standalone") - self.token_comm.token_back(self._smem_workspace, smem_base) + if cutlass.const_expr(not self._diag_skip_token_back): + self.token_comm.token_back(self._smem_workspace, smem_base) iket.range_pop() finalize_barrier.arrive_unaligned() self.token_comm.remove_device_members() diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/topk_reduce_dedup.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/topk_reduce_dedup.py new file mode 100644 index 000000000000..9e43d0dd1dda --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/topk_reduce_dedup.py @@ -0,0 +1,110 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # ctm-lint: allow-internal-marker +# SPDX-License-Identifier: BSD-3-Clause +"""Combine reduce for the dedup path: sums the ``rank_slots[token][rank]`` rows named by ``valid_rank_mask`` in +ascending rank order into ``reduced_output[token]``; top-k weights are already folded in.""" + +from typing import ClassVar + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cutlass_dsl import Float32, Int32 + +from .....helpers.dsl_helpers import mark_alignment +from .....quant_def import CombineFormat + + +class TopkReduceDedup: + _threads: ClassVar[int] = 128 + _hidden_per_thread: ClassVar[int] = 16 # one 32 B (ld/st.256) access of bf16 per slot + + def __init__(self, hidden: int, world_size: int, combine_format: CombineFormat) -> None: + if combine_format.is_quantized or combine_format.act_dtype is not cutlass.BFloat16: + raise ValueError("TopkReduceDedup supports the BF16 combine format only.") + if world_size > 32: + raise ValueError("TopkReduceDedup keeps the rank mask in an Int32; world_size must be <= 32.") + self.hidden = int(hidden) + self.world_size = int(world_size) + self.combine_format = combine_format + if self.hidden % self._hidden_per_thread != 0: + raise ValueError(f"hidden ({self.hidden}) must be divisible by {self._hidden_per_thread}.") + self.hidden_tiles = self.hidden // self._hidden_per_thread + self.require_predicate = self.hidden_tiles % self._threads != 0 + + @cute.jit + def __call__( + self, + rank_slots: cute.Tensor, # (token, world_size, hidden) bf16 + valid_rank_mask: cute.Tensor, # (token,) i32 bitmask over ranks + reduced_output: cute.Tensor, # (token, hidden) + stream: cuda.CUstream, + ): + threads = self._threads + total_workers = reduced_output.shape[0] * self.hidden_tiles + grid = [(total_workers + threads - 1) // threads, 1, 1] + block = [threads, 1, 1] + rank_slots = cute.make_tensor( + rank_slots.iterator, + cute.make_layout((rank_slots.shape[0], self.world_size, self.hidden), stride=rank_slots.stride), + ) + reduced_output = cute.make_tensor( + reduced_output.iterator, + cute.make_layout((reduced_output.shape[0], self.hidden), stride=reduced_output.stride), + ) + self._reduce_bf16(rank_slots, valid_rank_mask, reduced_output).launch(grid=grid, block=block, stream=stream) + + @cute.kernel + def _reduce_bf16(self, rank_slots: cute.Tensor, valid_rank_mask: cute.Tensor, reduced_output: cute.Tensor): + threads = self._threads + hidden_per_thread = self._hidden_per_thread + hidden_tiles = self.hidden_tiles + world_size: cutlass.Constexpr[int] = self.world_size + needs_guard = self.require_predicate + out_dtype = reduced_output.element_type + + worker_idx = cute.arch.block_idx()[0] * Int32(threads) + cute.arch.thread_idx()[0] + token_idx = worker_idx // hidden_tiles + hidden_tile_idx = worker_idx % hidden_tiles + + if (not needs_guard) or token_idx < reduced_output.shape[0]: + # (token, rank, hidden) -> (rank, hidden_per_thread) slice owned by this worker + slots = cute.zipped_divide(rank_slots[token_idx, None, None], (world_size, hidden_per_thread))[ + (None, None), (0, hidden_tile_idx) + ] + dst = cute.zipped_divide(reduced_output[token_idx, None], (hidden_per_thread,))[(None,), (hidden_tile_idx,)] + rank_mask = valid_rank_mask[token_idx] + + access_bits: cutlass.Constexpr[int] = hidden_per_thread * cutlass.BFloat16.width + load_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=access_bits) + acc = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float32) + for i in cutlass.range_constexpr(hidden_per_thread): + acc[i] = Float32(0.0) + one_pair = (Float32(1.0), Float32(1.0)) + + # Issue every rank slot's load first (predicated, one register buffer per rank) so all of them are in + # flight together, then accumulate ranks ascending; fma by 1.0 is an exact fp32 add. + terms = [cute.make_rmem_tensor((hidden_per_thread,), cutlass.BFloat16) for _ in range(world_size)] + for rank in cutlass.range_constexpr(world_size): + if ((rank_mask >> Int32(rank)) & Int32(1)) != Int32(0): + cute.copy( + load_atom, + mark_alignment(slots[rank, None], hidden_per_thread * cutlass.BFloat16.width // 8), + terms[rank], + ) + for rank in cutlass.range_constexpr(world_size): + if ((rank_mask >> Int32(rank)) & Int32(1)) != Int32(0): + for i in cutlass.range_constexpr(0, hidden_per_thread, 2): + value_pair = (Float32(terms[rank][i]), Float32(terms[rank][i + 1])) + acc[i], acc[i + 1] = cute.arch.fma_packed_f32x2(value_pair, one_pair, (acc[i], acc[i + 1])) + + out = cute.make_rmem_tensor((hidden_per_thread,), out_dtype) + out.store(acc.load().to(out_dtype)) + cute.copy( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), out_dtype, num_bits_per_copy=hidden_per_thread * out_dtype.width), + out, + mark_alignment(dst, hidden_per_thread * out_dtype.width // 8), + ) + + +__all__ = ["TopkReduceDedup"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/non_clc_mixed_cga.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/non_clc_mixed_cga.py index ba879cc6ac4d..cad226f2dd3a 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/non_clc_mixed_cga.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/non_clc_mixed_cga.py @@ -159,6 +159,11 @@ def __init__( # when some clusters only communicate. Atomic claims do not use this. self.work_id_stride = work_id_stride + @property + def max_fallback_cluster_count(self) -> int: + """Upper bound on physical fallback clusters: every preferred-cluster slot of the launch split.""" + return self.config.launch_cluster_cnt_merge_as_preferred * self.config.split_factor + def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: """Register the cluster-local claim broadcast channel.""" if self.work_id_mode == "atomic_counter": @@ -183,10 +188,12 @@ def register_device_workspace(self, device_workspace: DeviceWorkspace) -> None: buffer_space="local", reset="tail_reset", ) + # Under PDL more slots can split into fallback clusters than the host plan assumed: size the group + # table for every slot splitting. device_workspace.register( self.fallback_group_token_region, cutlass.Int64, - (self.config.fallback_cluster_count,), + (self.max_fallback_cluster_count,), buffer_space="local", reset="tail_reset", ) @@ -262,7 +269,7 @@ def assign_device_members( registration_counter_pointer=device_workspace.ptr(self.fallback_registration_counter_region), group_token_pointer=device_workspace.ptr(self.fallback_group_token_region), split_factor=self.config.split_factor, - fallback_cluster_count=self.config.fallback_cluster_count, + fallback_cluster_count=self.max_fallback_cluster_count, is_fallback_cluster=is_fallback_cluster, fallback_group_idx=Int32(0), in_group_idx=Int32(0), From 9b3b029973a3b5f2a39c5bd09808e9b206418b22 Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 04:48:53 -0700 Subject: [PATCH 08/59] [None][fix] Point the live-plane ABI contract test at moe/kernels The cute_dsl_kernels split on main moved megamoe_scheduler under tensorrt_llm/_torch/moe/kernels/, but test_vendored_live_plane_abi_matches_independent_goldens still loaded sami/geometry.py from the old cute_dsl_kernels path and failed with FileNotFoundError. Signed-off-by: jintaop --- tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index 1cb660a908e6..04a689f43468 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -344,7 +344,7 @@ def test_vendored_live_plane_abi_matches_independent_goldens( expected_stride_bytes, expected_planes, ): - path = _TORCH_ROOT / "cute_dsl_kernels/megamoe_scheduler/sami/geometry.py" + path = _TORCH_ROOT / "moe/kernels/megamoe_scheduler/sami/geometry.py" module_name = f"_live_plane_geometry_{hidden}_{intermediate}" spec = importlib.util.spec_from_file_location(module_name, path) assert spec is not None and spec.loader is not None From 313df46a7e622dc5bc46dec15ae87805e117dc1c Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 04:52:44 -0700 Subject: [PATCH 09/59] [None][perf] Cut per-layer host overhead of MoE rebalance and MegaMoE launches With per-layer dynamic EPLB on short-prompt prefill, rank 0 is host bound, so per-layer Python work on the rebalance and MegaMoE launch path costs iteration time directly. Outputs are unchanged. Rebalance submission: - HALO-Q and the TMA copy make their stream and device current through the raw torch._C setters (enter_cuda_stream / exit_cuda_stream) instead of the torch.cuda.device and torch.cuda.stream context managers. - RebalanceSlotSchedulerGroup._check_owner, called several times per layer, reads the raw current device and stream and returns early for the bound owner; it builds a torch.cuda.Stream only when binding. MegaMoE launch: - Eager launches call run_cute_dsl_megamoe_nvfp4_blackwell, the op's implementation, directly and skip the custom-op dispatch; under torch.compile the registered op is still used. - The backend owns a cache of the cute views of its weights and hands it to the op for the launch, like the profiling scratch. A hit requires an equal (data_ptr, shape, stride, dtype, device, alignment) key, so it is exact; the views keep their weights alive, so the cache lives and dies with the module. Tuning launches convert afresh. - Drop the second _tactic_cache_key computation in the runner's forward. Tests: CPU contract tests for the stream switch round trip, the owner fast path (other thread, stream or device still raise) and the weight view cache key. Signed-off-by: jintaop --- .../custom_ops/cute_dsl_megamoe_custom_op.py | 64 ++++++++--- .../fused_moe/mega_moe/mega_moe_cute_dsl.py | 17 ++- .../mega_moe/rebalance_slot_scheduler.py | 28 ++--- .../cuda_scheduler/runtime.py | 32 +++++- .../megamoe_scheduler/sami/hierarchical.py | 8 +- .../_torch/moe/test_dynamic_eplb_contracts.py | 104 +++++++++++++++++- 6 files changed, 213 insertions(+), 40 deletions(-) diff --git a/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py b/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py index fbe26c7c972f..de498b2cd01e 100644 --- a/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py +++ b/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py @@ -882,10 +882,7 @@ def _construct_megamoe_kernel( from cutlass.cute.nvgpu import OperandMajorMode from tensorrt_llm._torch.moe.kernels.cutedsl_megamoe.api import ImplDesc, ProblemDesc - from tensorrt_llm._torch.moe.kernels.cutedsl_megamoe.quant_def import ( - CombineFormat, - QuantKind, - ) + from tensorrt_llm._torch.moe.kernels.cutedsl_megamoe.quant_def import CombineFormat, QuantKind from tensorrt_llm._torch.moe.kernels.megamoe_kernels import ( BlackwellInferenceMegaMoE, RubinInferenceGenphaseMegaMoE, @@ -1554,6 +1551,19 @@ def to_cute_ptr(t, assumed_align=16): return to_cute, to_cute_ptr, SymBufferHost, cute + def _to_cute_cached(cache: dict, t: torch.Tensor, assumed_align: int = 16): + """Return the cute view of ``t`` from ``cache``, converting it on a miss. + + The key holds everything the view encodes, so a hit equals a fresh + conversion even if the address was reused. Only for tensors whose + layout is fixed (module weights): one entry per tensor. + """ + key = (t.data_ptr(), tuple(t.shape), t.stride(), t.dtype, t.device.index, assumed_align) + cute_tensor = cache.get(key) + if cute_tensor is None: + cute_tensor = cache[key] = _cute_launch_helpers()[0](t, assumed_align) + return cute_tensor + def _evict_latched_tuning_workspaces() -> None: """Pop every tuning-latched local workspace and prune its fence keys. @@ -1933,6 +1943,16 @@ def get_megamoe_profiling_scratch( _MEGAMOE_PROFILING_SCRATCH_CACHE[cache_key] = provider return provider + # Cute views of the calling module's weights, set by the backend around + # the op call like the profiling scratch. A cute view keeps its tensor + # alive, so the module owns the dict and the views never outlive it. + _ACTIVE_MEGAMOE_CUTE_WEIGHT_CACHE: Optional[dict] = None + + def set_active_megamoe_cute_weight_cache(cache: Optional[dict]) -> None: + """Set (``None`` clears) the calling module's cute weight-view cache.""" + global _ACTIVE_MEGAMOE_CUTE_WEIGHT_CACHE + _ACTIVE_MEGAMOE_CUTE_WEIGHT_CACHE = cache + def set_active_megamoe_profiling_scratch(regions) -> None: """Set (``None`` clears) the :class:`MegaMoeSymmRegions` used for AutoTuner profiling on the current call (set by the backend around @@ -2786,6 +2806,12 @@ def forward( # 32-bit memref shape field overflows once shared_workspace passes # 2 GiB (the kernel addresses by raw base + Int64 byte offset). _to_cute, _to_cute_ptr, SymBufferHost, cute = _cute_launch_helpers() + # Weights reuse the calling module's cute views; tuning launches convert afresh. + weight_cache = _ACTIVE_MEGAMOE_CUTE_WEIGHT_CACHE + if weight_cache is None or AutoTuner.get().is_tuning_mode: + _to_cute_w = _to_cute + else: + _to_cute_w = functools.partial(_to_cute_cached, weight_cache) # combine_output (max_T, 1, hidden) reshapes freely to 2D. Weights # present K stride-1 via a transpose VIEW (DLPack carries the @@ -2794,13 +2820,13 @@ def forward( runtime_kwargs = dict( activation=_to_cute(activation), activation_sf=_to_cute(activation_sf), - fc1_weight=_to_cute(fc1_weight.transpose(1, 2)), - fc1_weight_sf=_to_cute(fc1_weight_sf), - fc2_weight=_to_cute(fc2_weight.transpose(1, 2)), - fc2_weight_sf=_to_cute(fc2_weight_sf), - fc1_alpha=_to_cute(fc1_alpha, assumed_align=4), - fc2_alpha=_to_cute(fc2_alpha, assumed_align=4), - fc1_norm_const=_to_cute(fc1_norm_const, assumed_align=4), + fc1_weight=_to_cute_w(fc1_weight.transpose(1, 2)), + fc1_weight_sf=_to_cute_w(fc1_weight_sf), + fc2_weight=_to_cute_w(fc2_weight.transpose(1, 2)), + fc2_weight_sf=_to_cute_w(fc2_weight_sf), + fc1_alpha=_to_cute_w(fc1_alpha, assumed_align=4), + fc2_alpha=_to_cute_w(fc2_alpha, assumed_align=4), + fc1_norm_const=_to_cute_w(fc1_norm_const, assumed_align=4), output_activation=_to_cute(output_activation), local_workspace=_to_cute_ptr(local_workspace), shared_workspace=_to_cute_ptr(shared_workspace), @@ -2816,7 +2842,6 @@ def forward( max_ranks=int(self.world_size), ), ) - cache_key = self._tactic_cache_key(tactic_t) # READY ABI tail: attached only when S > 0, i.e. rebalance is on. # Contract: with S > 0 both values must be supplied. Fail loudly # here rather than hand the kernel an unarmed gate. @@ -2942,12 +2967,7 @@ def _make(): # ----- torch op --------------------------------------------------------- - @torch.library.custom_op( - "trtllm::cute_dsl_megamoe_nvfp4_blackwell", - mutates_args=("combine_output", "shared_workspace"), - device_types="cuda", - ) - def cute_dsl_megamoe_nvfp4_blackwell( + def run_cute_dsl_megamoe_nvfp4_blackwell( activation: torch.Tensor, activation_sf: torch.Tensor, topk_idx: torch.Tensor, @@ -3141,6 +3161,14 @@ def cute_dsl_megamoe_nvfp4_blackwell( shared_workspace=shared_workspace, ) + # Eager callers may call ``run_cute_dsl_megamoe_nvfp4_blackwell`` directly + # and skip the custom-op dispatch; traced code must use the registered op. + cute_dsl_megamoe_nvfp4_blackwell = torch.library.custom_op( + "trtllm::cute_dsl_megamoe_nvfp4_blackwell", + mutates_args=("combine_output", "shared_workspace"), + device_types="cuda", + )(run_cute_dsl_megamoe_nvfp4_blackwell) + @torch.library.register_fake("trtllm::cute_dsl_megamoe_nvfp4_blackwell") def _( activation: torch.Tensor, diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py index 04fb4ef0723c..f6d6663eb98f 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py @@ -72,7 +72,7 @@ import socket import weakref from dataclasses import dataclass -from typing import Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Tuple, Union import torch import torch.distributed as dist @@ -761,6 +761,8 @@ def __init__( # deferred factory in ``_launch_megamoe_kernel``. self._profiling_scratch_provider_ref = None self._profiling_scratch_kwargs = None + # Cute views of this module's weights, reused across launches. + self._cute_weight_cache: Dict[tuple, Any] = {} self._weights_loaded = False self._weights_created = False self.quant_method = None @@ -1787,6 +1789,8 @@ def _launch_megamoe_kernel( # a later same-process call for a different shape never sees a stale # scratch (ownership: see ``_profiling_scratch_provider_ref``). from tensorrt_llm._torch.moe.custom_ops.cute_dsl_megamoe_custom_op import ( + run_cute_dsl_megamoe_nvfp4_blackwell, + set_active_megamoe_cute_weight_cache, set_active_megamoe_profiling_scratch, set_active_megamoe_profiling_scratch_factory, ) @@ -1851,8 +1855,16 @@ def _deferred_scratch_factory(_self=self, _get=get_megamoe_profiling_scratch): # Quantization, independent staging, output clears, and host setup # are already enqueued. The router inside this op first reads routes. self._wait_rebalance_routes(replica_plan) + # Eager launches call the op's implementation directly: the custom-op + # dispatch costs tens of microseconds of host time per layer. + megamoe_op = ( + torch.ops.trtllm.cute_dsl_megamoe_nvfp4_blackwell + if torch.compiler.is_compiling() + else run_cute_dsl_megamoe_nvfp4_blackwell + ) + set_active_megamoe_cute_weight_cache(self._cute_weight_cache) try: - torch.ops.trtllm.cute_dsl_megamoe_nvfp4_blackwell( + megamoe_op( activation=_as_nvfp4(activation), activation_sf=_as_fp8_sf(activation_sf), topk_idx=topk_idx, @@ -1902,6 +1914,7 @@ def _deferred_scratch_factory(_self=self, _get=get_megamoe_profiling_scratch): # global alive (its closure strongly references this module). set_active_megamoe_profiling_scratch(None) set_active_megamoe_profiling_scratch_factory(None) + set_active_megamoe_cute_weight_cache(None) if num_tokens == 0: return torch.empty((0, hidden), dtype=output_dtype, device=combine_output.device) # The kernel already collapsed the top-k axis into the (T, 1, hidden) diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py index c603c55125d9..ba46b7025b05 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py @@ -251,22 +251,22 @@ def _check_owner(self) -> None: if self._closed: raise RuntimeError("Rebalance scheduler group is closed") thread_id = threading.get_ident() - if torch.cuda.current_device() != self.device: + # Runs several times per layer: read the raw device and stream instead + # of building a torch.cuda.Stream, and return early for the bound owner. + if torch._C._cuda_getDevice() != self.device: raise RuntimeError("Rebalance submission requires its bound CUDA device") - stream = torch.cuda.current_stream(self.device) - handle = int(stream.cuda_stream) - if self._owner_thread_id is None: - if ( - self._execution_stream_handle is not None - and handle != self._execution_stream_handle - ): - raise RuntimeError("Rebalance MAIN stream cannot change during owner handoff") - if handle == self._stream_handle or stream.priority <= self.copy_stream.priority: - raise RuntimeError("MAIN must be distinct from the higher-priority COPY stream") - self._owner_thread_id = thread_id - self._execution_stream_handle = handle - elif thread_id != self._owner_thread_id or handle != self._execution_stream_handle: + handle = torch._C._cuda_getCurrentRawStream(self.device) + if thread_id == self._owner_thread_id and handle == self._execution_stream_handle: + return + if self._owner_thread_id is not None: raise RuntimeError("Rebalance must be submitted by one MAIN thread and stream") + if self._execution_stream_handle is not None and handle != self._execution_stream_handle: + raise RuntimeError("Rebalance MAIN stream cannot change during owner handoff") + stream = torch.cuda.current_stream(self.device) + if handle == self._stream_handle or stream.priority <= self.copy_stream.priority: + raise RuntimeError("MAIN must be distinct from the higher-priority COPY stream") + self._owner_thread_id = thread_id + self._execution_stream_handle = handle def release_warmup_owner_after_sync(self, execution_stream_handle: int) -> None: """Release the startup CPU owner after the caller drains this device. diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py index b88af7be64ba..cd6ea2d9491b 100644 --- a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py @@ -199,6 +199,33 @@ def __init__(self, config: CudaSchedulerConfig, device: torch.device): ) +def enter_cuda_stream(stream: torch.cuda.Stream) -> tuple: + """Make ``stream`` and its device current; ``exit_cuda_stream`` restores both. + + Same effect as entering ``torch.cuda.device(stream.device)`` and then + ``torch.cuda.stream(stream)``, without their Python overhead on the + per-layer submission path. Returns the state to pass to ``exit_cuda_stream``. + """ + index = stream.device_index + prev_device = torch._C._cuda_exchangeDevice(index) + prev_stream = torch._C._cuda_getCurrentStream(index) + torch._C._cuda_setStream( + stream_id=stream.stream_id, + device_index=index, + device_type=stream.device_type, + ) + return prev_device, prev_stream + + +def exit_cuda_stream(state: tuple) -> None: + """Restore the stream and device saved by ``enter_cuda_stream``.""" + prev_device, (stream_id, device_index, device_type) = state + torch._C._cuda_setStream( + stream_id=stream_id, device_index=device_index, device_type=device_type + ) + torch._C._cuda_maybeExchangeDevice(prev_device) + + def _stream_handle(stream: object) -> int: if isinstance(stream, bool): raise TypeError("stream must be a CUDA stream handle") @@ -455,7 +482,8 @@ def _launch_op( stream = self._plan_channel_torch_stream if stream is None or _stream_handle(stream) != stream_handle: stream = torch.cuda.ExternalStream(stream_handle, device=self.device) - with torch.cuda.device(self.device), torch.cuda.stream(stream): + state = enter_cuda_stream(stream) + try: torch.ops.trtllm.moe_rebalance_halo_q( routes, self.outputs.physical_slot_ids, @@ -484,6 +512,8 @@ def _launch_op( self._plan_channel_route_features, valid_route_count, ) + finally: + exit_cuda_stream(state) def submit( self, routes: torch.Tensor, stream_handle: int, *, diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py index b6915a6fd3a7..8669b24d00a2 100644 --- a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py @@ -14,6 +14,7 @@ from .arena import BoundLiveWeightArena, LiveWeightArena, _TensorIdentity from ._util import check_cuda +from ..cuda_scheduler.runtime import enter_cuda_stream, exit_cuda_stream from ..geometry import hierarchy_group_sizes from .geometry import BundleLayout @@ -634,9 +635,8 @@ def submit( raise RuntimeError("outputs are not from the bound CUDA scheduler") self._require_live_bank_reusable() try: - with torch.cuda.device(self.device), torch.cuda.stream( - self._stream_owner - ): + state = enter_cuda_stream(self._stream_owner) + try: if start_event is not None: start_event.record() commands = int( @@ -648,6 +648,8 @@ def submit( payload_end_event.record() if end_event is not None: end_event.record() + finally: + exit_cuda_stream(state) self._active_generation = self._generation release = getattr(self._scheduler, "release_plan_channel") release() diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index 04a689f43468..84e6e588a65b 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -66,7 +66,7 @@ def _extract_class( def _extract_function(path: Path, name: str, namespace: dict): tree = ast.parse(path.read_text()) function = next( - node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == name + node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == name ) function.decorator_list = [] module = ast.Module( @@ -1066,7 +1066,11 @@ def test_direct_submit_orders_generations_and_defers_the_route_wait(): current_stream=lambda _: main, nvtx=SimpleNamespace(range=lambda _: nullcontext()), ) - torch = SimpleNamespace(cuda=cuda, Tensor=_Routes, int32="int32") + raw = SimpleNamespace( + _cuda_getDevice=lambda: state.device, + _cuda_getCurrentRawStream=lambda _: main.cuda_stream, + ) + torch = SimpleNamespace(cuda=cuda, _C=raw, Tensor=_Routes, int32="int32") source = _TORCH_ROOT / "moe/fused_moe/mega_moe/rebalance_slot_scheduler.py" namespace = {"torch": torch, "threading": SimpleNamespace(get_ident=lambda: state.thread)} group_cls = _extract_class(source, "RebalanceSlotSchedulerGroup", None, namespace) @@ -1148,6 +1152,20 @@ def copy_submit(result): assert trace.index(("TMA", generation)) < trace.index(("wait", 44, 2)) assert trace[-2:] == [("record", 3, 44), ("consumer_release", generation, 3)] + # The bound-owner fast path still rejects another thread, stream or device. + state.thread = 8 + with pytest.raises(RuntimeError, match="one MAIN thread and stream"): + group.plan_schedule(_Routes(8, trace)) + state.thread = 7 + main.cuda_stream = 45 + with pytest.raises(RuntimeError, match="one MAIN thread and stream"): + group.plan_schedule(_Routes(8, trace)) + main.cuda_stream = 44 + state.device = 1 + with pytest.raises(RuntimeError, match="bound CUDA device"): + group.plan_schedule(_Routes(8, trace)) + state.device = 0 + trace.clear() with pytest.raises(ValueError, match="contiguous CUDA int32"): group.plan_schedule(_Routes(8, trace, dtype="int64")) @@ -1160,6 +1178,88 @@ def copy_submit(result): assert group._plan_part is None +def test_raw_stream_switch_restores_the_previous_stream_and_device(): + calls = [] + current = {"device": 0, "streams": {0: (10, 0, 1), 3: (30, 3, 1)}} + + def exchange_device(index): + calls.append(("exchange", index)) + previous, current["device"] = current["device"], index + return previous + + def set_stream(stream_id, device_index, device_type): + calls.append(("set", stream_id, device_index)) + current["streams"][device_index] = (stream_id, device_index, device_type) + + def maybe_exchange_device(index): + calls.append(("restore_device", index)) + current["device"] = index + return index + + raw = SimpleNamespace( + _cuda_exchangeDevice=exchange_device, + _cuda_getCurrentStream=lambda index: current["streams"][index], + _cuda_setStream=set_stream, + _cuda_maybeExchangeDevice=maybe_exchange_device, + ) + source = _TORCH_ROOT / "moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py" + namespace = {"torch": SimpleNamespace(_C=raw, cuda=SimpleNamespace(Stream=object))} + enter = _extract_function(source, "enter_cuda_stream", namespace) + leave = _extract_function(source, "exit_cuda_stream", namespace) + + state = enter(SimpleNamespace(device_index=3, stream_id=77, device_type=1)) + assert current["device"] == 3 and current["streams"][3] == (77, 3, 1) + leave(state) + assert current == {"device": 0, "streams": {0: (10, 0, 1), 3: (30, 3, 1)}} + assert calls == [ + ("exchange", 3), + ("set", 77, 3), + ("set", 30, 3), + ("restore_device", 0), + ] + + +def test_megamoe_cute_weight_cache_keys_every_view_field(): + conversions = [] + + def to_cute(t, assumed_align): + conversions.append((t, assumed_align)) + return ("cute", len(conversions)) + + to_cute_cached = _extract_function( + _TORCH_ROOT / "moe/custom_ops/cute_dsl_megamoe_custom_op.py", + "_to_cute_cached", + {"torch": SimpleNamespace(Tensor=object), "_cute_launch_helpers": lambda: (to_cute,)}, + ) + + def weight(ptr=4096, shape=(8, 16), stride=(16, 1), dtype="fp4", device=0): + return SimpleNamespace( + data_ptr=lambda: ptr, + shape=shape, + stride=lambda: stride, + dtype=dtype, + device=SimpleNamespace(index=device), + ) + + cache = {} + first = to_cute_cached(cache, weight()) + # A new view object of the same memory and layout reuses the cute view. + assert to_cute_cached(cache, weight()) is first + assert len(conversions) == 1 + for changed in ( + weight(ptr=8192), + weight(shape=(16, 8)), + weight(stride=(1, 8)), + weight(dtype="fp8"), + weight(device=1), + ): + assert to_cute_cached(cache, changed) is not first + assert to_cute_cached(cache, weight(), assumed_align=4) is not first + assert len(conversions) == 7 + # The cache belongs to the caller: another module's dict starts empty. + assert to_cute_cached({}, weight()) is not first + + def test_route_wait_occurs_after_independent_staging(): source = _TORCH_ROOT / "moe/fused_moe/mega_moe/mega_moe_cute_dsl.py" impl = _extract_class( From ace269ae82d0f18794f4a1f296396b77f3f626f1 Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 04:51:46 -0700 Subject: [PATCH 10/59] [None][feat] Add ENABLE_PERFECT_ROUTER=expert_ids to bound the MoE load-balance gain ENABLE_PERFECT_ROUTER=1 replaces the router logits with synthesized ones. That cannot steer routers whose expert ids do not come from the logits (the DeepSeek-V4 hash layers), and the DeepSeekV3 planner that DeepSeek-V4 uses plans each new token count on the host, which a context-only run with a different token count every iteration cannot afford. ENABLE_PERFECT_ROUTER=expert_ids keeps the learned router and its routing method, then replaces the routed expert ids in the MoE scheduler with the same balanced assignment that the flat perfect-router planner targets: slot t * top_k + j, rotated by the EP rank, goes to rank slot % EP and to expert (slot // EP) % (E / EP) on it. Every rank and expert receives an equal share of the routes, within one, for any token count. The ids are a view of one cached int32 table per expert geometry and device, built on first use and grown by doubling, so a call launches no kernel. Routing scales are left as computed; outputs are wrong by design. Both scheduler routing sites apply it before the load balancer and route capture see the ids, so it covers every ConfigurableMoE backend that routes outside its kernel, including MEGAMOE_CUTEDSL. Run with the load balancer off, the gap to a plain run bounds what any MoE load balancer can gain. Default unchanged. Adds a CPU test of the table and documents the mode. Signed-off-by: jintaop --- docs/source/developer-guide/perf-analysis.md | 27 ++++++++ .../_torch/moe/fused_moe/interface.py | 30 +++++++- .../_torch/moe/fused_moe/moe_scheduler.py | 8 ++- tensorrt_llm/_torch/moe/fused_moe/routing.py | 54 +++++++++++++++ tests/unittest/_torch/moe/test_moe_routing.py | 68 ++++++++++++++++++- 5 files changed, 181 insertions(+), 6 deletions(-) diff --git a/docs/source/developer-guide/perf-analysis.md b/docs/source/developer-guide/perf-analysis.md index c0dcebe7b94d..8e4830b51a32 100644 --- a/docs/source/developer-guide/perf-analysis.md +++ b/docs/source/developer-guide/perf-analysis.md @@ -237,6 +237,33 @@ ENABLE_PERFECT_ROUTER=1 trtllm-serve ... | Similar performance with/without perfect router | Router load balancing is not a bottleneck; focus optimization efforts elsewhere | | Significant improvement with perfect router | The learned router is causing load imbalance; consider router optimization or load balancing strategies | +### Replacing Expert IDs Instead of Logits + +`ENABLE_PERFECT_ROUTER=expert_ids` keeps the learned router and its routing +method, then replaces the routed expert IDs with a perfectly balanced +assignment: every EP rank and every expert receives an equal share of the +routes (within one) for any token count, and the receiver order is rotated by +the sender rank. The routing scales are left as computed. + +Use it where logits cannot steer the routing, for example the DeepSeek-V4 hash +layers, which pick experts from the token IDs, or where the token count changes +every iteration, as in context-only benchmarks: the IDs come from one cached +table per expert geometry, so no logits are synthesized per token count. + +It applies to MoE backends that route through `ConfigurableMoE`, including +`MEGAMOE_CUTEDSL`. With the load balancer off, the gap to a plain run bounds +what any MoE load balancer can gain: + +```bash +# Learned routing, load balancer off +trtllm-bench ... +# Perfectly balanced routing, load balancer off (upper bound) +ENABLE_PERFECT_ROUTER=expert_ids trtllm-bench ... +``` + +The layer-wise benchmarks (`examples/layer_wise_benchmarks`) accept the same +environment variable with `--balance-method NotModified`. + ### Supported Models ```{note} diff --git a/tensorrt_llm/_torch/moe/fused_moe/interface.py b/tensorrt_llm/_torch/moe/fused_moe/interface.py index 074a7d026b78..6291e67d5a5e 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/interface.py +++ b/tensorrt_llm/_torch/moe/fused_moe/interface.py @@ -53,6 +53,7 @@ def _reject(reason: MoERejectReason, detail: str) -> MoEEligibility: get_model_extra_attrs, is_gated_activation, is_torch_compiling) from .routing import (BaseMoeRoutingMethod, RoutingMethodType, + get_cached_perfect_router_expert_ids, get_cached_perfect_router_logits, precompute_common_perfect_router_logits) @@ -466,8 +467,12 @@ def _get_perfect_router_dtype(self) -> torch.dtype: return self.dtype if self.dtype is not None else torch.float32 def _init_perfect_router(self): - self._enable_perfect_router = os.environ.get("ENABLE_PERFECT_ROUTER", - "0") == "1" + # ENABLE_PERFECT_ROUTER=1 replaces the router logits before routing; + # ENABLE_PERFECT_ROUTER=expert_ids replaces the routed expert ids after + # routing. Both are for performance analysis only. + perfect_router = os.environ.get("ENABLE_PERFECT_ROUTER", "0") + self._enable_perfect_router = perfect_router == "1" + self._enable_perfect_router_expert_ids = perfect_router == "expert_ids" if not self._enable_perfect_router: return @@ -496,6 +501,27 @@ def _maybe_get_perfect_router_logits( dtype=router_logits.dtype, routing_method=self.routing_method) + def _maybe_get_perfect_router_expert_ids( + self, token_selected_experts: torch.Tensor) -> torch.Tensor: + """Return balanced expert ids under ENABLE_PERFECT_ROUTER=expert_ids. + + The router still runs, so its cost is kept; only its expert ids are + replaced, by a view of a cached table (see + ``get_cached_perfect_router_expert_ids``). The routing scales are left + as computed. Model outputs are wrong by design. + """ + if not self._enable_perfect_router_expert_ids: + return token_selected_experts + + num_tokens, experts_per_token = token_selected_experts.shape + return get_cached_perfect_router_expert_ids( + num_tokens=num_tokens, + num_experts=self.num_experts, + experts_per_token=experts_per_token, + moe_ep_size=self.ep_size, + ep_rank=self.ep_rank, + device=token_selected_experts.device) + def _init_load_balancer( self, model_config: ModelConfig, diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py index f97969c566e9..2076be15b560 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py @@ -491,7 +491,9 @@ def _forward_chunk_impl( token_selected_experts, token_final_scales, x, x_sf = fused_result used_fused_route_quant = True - token_selected_experts = token_selected_experts.to(torch.int32) + token_selected_experts = moe._maybe_get_perfect_router_expert_ids( + token_selected_experts.to(torch.int32) + ) if row_offset is not None: route_capture = get_active_route_capture() if route_capture is not None: # R3: in-graph device-buffer capture @@ -1498,7 +1500,9 @@ def _forward_chunk( token_selected_experts, token_final_scales = moe.routing_method.apply_with_aux( router_logits_chunk_real, input_ids_chunk_real, routing_aux[:num_tokens] ) - token_selected_experts = token_selected_experts.to(torch.int32) + token_selected_experts = moe._maybe_get_perfect_router_expert_ids( + token_selected_experts.to(torch.int32) + ) route_capture = get_active_route_capture() if route_capture is not None: # R3: in-graph device-buffer capture route_capture.capture(moe.layer_idx, token_selected_experts, row_offset) diff --git a/tensorrt_llm/_torch/moe/fused_moe/routing.py b/tensorrt_llm/_torch/moe/fused_moe/routing.py index 432f93c4557a..f9b9c86201ce 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/routing.py +++ b/tensorrt_llm/_torch/moe/fused_moe/routing.py @@ -15,6 +15,12 @@ # Global cache for perfect router logits shared across MoE modules _PERFECT_ROUTER_LOGITS_CACHE = {} +# Global cache for perfect router expert-id tables shared across MoE modules, +# keyed by expert geometry, sender rank and device +_PERFECT_ROUTER_EXPERT_IDS_CACHE = {} +# Initial expert-id table rows, enough for the per-rank token counts of typical +# prefill chunks, so the table is normally built once, during warmup +_PERFECT_ROUTER_EXPERT_IDS_MIN_ROWS = 16384 def _get_perfect_router_cache_signature(routing_method): @@ -234,6 +240,54 @@ def get_cached_perfect_router_logits(num_tokens: int, num_experts: int, return logits +def get_cached_perfect_router_expert_ids(num_tokens: int, num_experts: int, + experts_per_token: int, + moe_ep_size: int, ep_rank: int, + device: torch.device) -> torch.Tensor: + """ + Return perfectly balanced expert ids for ``num_tokens`` local tokens. + + This is the ``ENABLE_PERFECT_ROUTER=expert_ids`` counterpart of + ``get_cached_perfect_router_logits``: instead of synthesizing logits for + the routing method to consume, it replaces the routed expert ids after + routing. It therefore also covers routers whose ids do not come from the + logits (DeepSeek-V4 hash layers) and needs no per-token-count planning. + + The ids follow ``FlatTopKPerfectRouterPlanner._get_balanced_expert_indices``: + slot ``t * experts_per_token + j``, rotated by ``ep_rank``, goes to EP rank + ``slot % moe_ep_size`` and to expert ``(slot // moe_ep_size) % + (num_experts // moe_ep_size)`` on it. A slot's expert depends only on its + own index, so the ids for ``n`` tokens are the leading ``n`` rows of any + longer table. One table per geometry and device is built on first use and + grown by doubling; every call returns a view of its leading rows, with no + kernel launch. + + Args: + num_tokens: Number of local tokens on the current EP rank. + num_experts: Total number of experts across all EP ranks. + experts_per_token: Number of experts selected for each token. + moe_ep_size: Number of EP ranks. + ep_rank: Sender EP rank used to rotate the assignment stream. + device: Device of the routed expert ids. + + Returns: + torch.Tensor: Contiguous int32 expert ids of shape + ``[num_tokens, experts_per_token]``. + """ + cache_key = (num_experts, experts_per_token, moe_ep_size, ep_rank, device) + table = _PERFECT_ROUTER_EXPERT_IDS_CACHE.get(cache_key) + if table is None or table.shape[0] < num_tokens: + if table is None: + num_rows = _PERFECT_ROUTER_EXPERT_IDS_MIN_ROWS + else: + num_rows = 2 * table.shape[0] + table = FlatTopKPerfectRouterPlanner._get_balanced_expert_indices( + max(num_rows, num_tokens), num_experts, experts_per_token, + moe_ep_size, ep_rank, device).to(torch.int32) + _PERFECT_ROUTER_EXPERT_IDS_CACHE[cache_key] = table + return table[:num_tokens] + + # The type of method in top-K routing, for use in torch custom op # Please keep this in sync with the counterpart defined in cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.h class RoutingMethodType(IntEnum): diff --git a/tests/unittest/_torch/moe/test_moe_routing.py b/tests/unittest/_torch/moe/test_moe_routing.py index 31503d512dbd..1c286264cffd 100644 --- a/tests/unittest/_torch/moe/test_moe_routing.py +++ b/tests/unittest/_torch/moe/test_moe_routing.py @@ -22,8 +22,8 @@ RenormalizeNaiveMoeRoutingMethod, SparseMixerMoeRoutingMethod, StaticMoeRoutingMethod, create_load_balanced_logits, create_moe) from tensorrt_llm._torch.moe.fused_moe import routing as moe_routing -from tensorrt_llm._torch.moe.fused_moe.routing import \ - get_cached_perfect_router_logits +from tensorrt_llm._torch.moe.fused_moe.routing import ( + get_cached_perfect_router_expert_ids, get_cached_perfect_router_logits) from tensorrt_llm._torch.nccl_window_tensor_scope import \ discard_nccl_window_tensor_outputs from tensorrt_llm._utils import mpi_rank @@ -636,6 +636,70 @@ def test_perfect_router_get_cached_logits_e2e(routing_name, shape): assert torch.equal(logits, logits_2) +@pytest.mark.parametrize( + "num_experts,top_k,moe_ep_size", + [ + pytest.param(64, 6, 4, id="small"), + pytest.param(384, 6, 8, id="dsv4_pro_ep8"), + ], +) +def test_perfect_router_expert_ids_balanced(num_experts, top_k, moe_ep_size): + """ENABLE_PERFECT_ROUTER=expert_ids: balanced ids for any token count. + + Checks, on CPU, for every sender rank and a range of token counts: + - every EP rank and every expert receives an equal share of the routes, + within one; + - a token's routes land on distinct EP ranks and distinct experts; + - the ids for fewer tokens are a view of the leading rows of the cached + table, also after the table grows. + """ + device = torch.device("cpu") + experts_per_rank = num_experts // moe_ep_size + moe_routing._PERFECT_ROUTER_EXPERT_IDS_CACHE.clear() + + for ep_rank in range(moe_ep_size): + largest = None + for num_tokens in [1, 3, 7, 100, 1000, 8191]: + ids = get_cached_perfect_router_expert_ids(num_tokens=num_tokens, + num_experts=num_experts, + experts_per_token=top_k, + moe_ep_size=moe_ep_size, + ep_rank=ep_rank, + device=device) + assert ids.shape == (num_tokens, top_k) + assert ids.dtype == torch.int32 + assert ids.is_contiguous() + + rank_counts = torch.bincount(ids.flatten() // experts_per_rank, + minlength=moe_ep_size) + assert rank_counts.max() - rank_counts.min() <= 1 + expert_counts = torch.bincount(ids.flatten(), minlength=num_experts) + assert expert_counts.max() - expert_counts.min() <= 1 + + if top_k <= moe_ep_size: + ranks = ids // experts_per_rank + assert all(len(set(row)) == top_k for row in ranks.tolist()) + assert all(len(set(row)) == top_k for row in ids.tolist()) + largest = ids + + # Growing the table keeps the leading rows, so cached views stay valid. + grown = get_cached_perfect_router_expert_ids( + num_tokens=2 * moe_routing._PERFECT_ROUTER_EXPERT_IDS_MIN_ROWS + 1, + num_experts=num_experts, + experts_per_token=top_k, + moe_ep_size=moe_ep_size, + ep_rank=ep_rank, + device=device) + assert torch.equal(grown[:largest.shape[0]], largest) + again = get_cached_perfect_router_expert_ids(num_tokens=5, + num_experts=num_experts, + experts_per_token=top_k, + moe_ep_size=moe_ep_size, + ep_rank=ep_rank, + device=device) + assert again.data_ptr() == grown.data_ptr() + + # ----------------------------------------------------------------- # Unified routing-method builder. Covers every RoutingMethodType in # tensorrt_llm._torch.moe.fused_moe.routing: From 7e03243d778c9629d9350f298eeaafbc296be842 Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 04:52:26 -0700 Subject: [PATCH 11/59] [None][feat] Layer-wise benchmarks: MoE load balancer and mid-model DeepSeek-V4 slices Let the layer-wise runner profile MegaMoE layers with per-layer dynamic EPLB beside the plain run and the ENABLE_PERFECT_ROUTER=expert_ids bound, including DeepSeek-V4 slices that start mid-model: - run.py --moe-load-balancer / Runner(moe_load_balancer=...) set MoeConfig.load_balancer: a YAML file path, or a mapping in the config file. The runner keeps the load balancer the model loader creates, runs every forward inside its iteration context, as the executor does, and shuts it down at the end of run.py. - Eager runs warm up on the stream they are timed on. Per-layer dynamic EPLB binds its submission stream on first use (the autotune pass) and rejects any other, so warming up on the CUDA-graph capture stream failed. CUDA-graph runs keep the capture-stream warmup. - DeepSeek-V4 decoder layers pass an mHC HCState between layers, and only layer 0 takes the plain residual, so a slice such as --layer-indices 7,8 failed. Such a slice now enters with a resolved HCState. Slices starting at layer 0 and models without mHC are unchanged. - README: DeepSeek-V4 MEGAMOE_CUTEDSL runs for the three arms. Signed-off-by: jintaop --- examples/layer_wise_benchmarks/README.md | 14 +++++++++ examples/layer_wise_benchmarks/run.py | 19 ++++++++++-- .../tools/layer_wise_benchmarks/runner.py | 29 +++++++++++++++++-- 3 files changed, 56 insertions(+), 6 deletions(-) diff --git a/examples/layer_wise_benchmarks/README.md b/examples/layer_wise_benchmarks/README.md index 5da3265d3eec..7bb0668f6164 100644 --- a/examples/layer_wise_benchmarks/README.md +++ b/examples/layer_wise_benchmarks/README.md @@ -121,6 +121,20 @@ NP=4 ./mpi_launch.sh ./run.sh config_gen.yaml --model $K3 --layer-indices 4,5,6, # Run Kimi K3 at the serving parallelism (DEP16, see examples/kimi_k3/eval_extra_llm_options.yaml) NP=16 ./mpi_launch.sh ./run.sh config_gen.yaml --model $K3 --layer-indices 4,5,6,7 --tokens-per-block 64 --moe-backend TRTLLM --seq-len-q 1 --moe-max-num-tokens 33024 --use-low-precision-moe-combine +# Run DeepSeek-V4 with MegaMoE, $DSV4 being a local checkpoint directory, as three arms: +# load balancer off, per-layer dynamic EPLB (num_slots = experts + helper slots per rank * EP, +# here 384 + 3 * 8), and perfectly balanced routing with the load balancer off (the upper +# bound for any MoE load balancer). MEGAMOE_CUTEDSL is not in the balance-method whitelist, +# so the arms use --balance-method NotModified. To replay the routing of an end-to-end run, +# add --replay-file-path (see "Performance alignment" below); the third arm replaces the +# routed expert IDs either way. A slice may start mid-model (here 7,8). +NP=8 ./mpi_launch.sh ./run.sh config_ctx.yaml --model $DSV4 --layer-indices 7,8 --moe-backend MEGAMOE_CUTEDSL --balance-method NotModified +NP=8 ./mpi_launch.sh ./run.sh config_ctx.yaml --model $DSV4 --layer-indices 7,8 --moe-backend MEGAMOE_CUTEDSL --balance-method NotModified --moe-load-balancer eplb_per_layer_dynamic.yaml +NP=8 ./mpi_launch.sh -x ENABLE_PERFECT_ROUTER=expert_ids ./run.sh config_ctx.yaml --model $DSV4 --layer-indices 7,8 --moe-backend MEGAMOE_CUTEDSL --balance-method NotModified +# with eplb_per_layer_dynamic.yaml holding: +# mode: per_layer_dynamic +# num_slots: 408 + # Run with DeepEP A2A NP=4 ./mpi_launch.sh -x TRTLLM_FORCE_COMM_METHOD=DEEPEP ./run.sh config_ctx.yaml --moe-backend CUTEDSL NP=4 ./mpi_launch.sh -x TRTLLM_FORCE_COMM_METHOD=DEEPEP ./run.sh config_gen.yaml --moe-backend CUTEDSL diff --git a/examples/layer_wise_benchmarks/run.py b/examples/layer_wise_benchmarks/run.py index 8d4a17b1ea1d..476c6ce18114 100644 --- a/examples/layer_wise_benchmarks/run.py +++ b/examples/layer_wise_benchmarks/run.py @@ -76,6 +76,12 @@ def comma_separated_floats(s): # whose Literal is the authoritative list. A second list here can only drift. parser.add_argument("--moe-backend-for-prefill", type=str) parser.add_argument("--moe-max-num-tokens", type=int) +parser.add_argument( + "--moe-load-balancer", + type=str, + help="MoeConfig.load_balancer: a YAML file path, or a mapping when set in the config" + " file, e.g. {mode: per_layer_dynamic, num_slots: 408}. Default: no load balancer.", +) parser.add_argument( "--spec-max-draft-len", type=int, @@ -252,6 +258,7 @@ def comma_separated_floats(s): use_cuda_graph=args.use_cuda_graph, spec_config=spec_config, vision_config=args.vision_config, + moe_load_balancer=args.moe_load_balancer, ) logger.info("Layer-wise benchmarks: Create runner ... Done") @@ -420,10 +427,15 @@ def comma_separated_floats(s): with runner.replace_routing_method_ctx( balance_method=BalanceMethod[args.balance_method], balance_ratio=balance_ratio ): - capture_stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(capture_stream): + if args.use_cuda_graph: + capture_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(capture_stream): + run_pack(check=True) + torch.cuda.current_stream().wait_stream(capture_stream) + else: + # Eager runs warm up on the stream they are timed on: per-layer + # dynamic EPLB binds its submission stream on first use. run_pack(check=True) - torch.cuda.current_stream().wait_stream(capture_stream) torch.cuda.synchronize() logger.info("Layer-wise benchmarks: Warmup ... Done") @@ -509,3 +521,4 @@ def comma_separated_floats(s): f" (us)" ) torch.cuda.cudart().cudaProfilerStop() +runner.shutdown_moe_load_balancer() diff --git a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py index 4382827aad45..32ffb5893e1f 100644 --- a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py +++ b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py @@ -9,7 +9,7 @@ import weakref from dataclasses import replace from enum import IntEnum -from typing import Literal, Optional +from typing import Any, Dict, Literal, Optional, Union import torch @@ -22,7 +22,12 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_utils import PostInitCaller, remove_weights, skip_forward from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata +from tensorrt_llm._torch.modules.mhc.hyper_connection import HCState from tensorrt_llm._torch.moe.fused_moe.fused_moe_trtllm_gen import TRTLLMGenFusedMoE +from tensorrt_llm._torch.moe.fused_moe.moe_load_balancer import ( + BaseMoeLoadBalancer, + MoeLoadBalancerIterContext, +) from tensorrt_llm._torch.nccl_window_tensor_scope import discard_nccl_window_tensor_outputs from tensorrt_llm._torch.pyexecutor._util import _mamba_conv_layout_kwargs, get_kv_cache_manager_cls from tensorrt_llm._torch.pyexecutor.config_utils import ( @@ -411,6 +416,7 @@ def __init__( use_cuda_graph: bool, spec_config: Optional[DecodingBaseConfig] = None, vision_config: Optional[str] = None, + moe_load_balancer: Optional[Union[str, Dict[str, Any]]] = None, ) -> None: super().__init__() @@ -438,6 +444,7 @@ def __init__( max_num_tokens=moe_max_num_tokens, disable_finalize_fusion=False, use_low_precision_moe_combine=use_low_precision_moe_combine, + load_balancer=moe_load_balancer, ), attn_backend="TRTLLM", kv_cache_config=KvCacheConfig( @@ -464,7 +471,7 @@ def __init__( max_num_tokens=max_num_tokens, max_seq_len=max_seq_len, ) - model, _ = model_loader.load( + model, load_balancer = model_loader.load( checkpoint_dir=pretrained_model_name_or_path, checkpoint_loader=checkpoint_loader ) @@ -520,6 +527,11 @@ def forward_block_residual( self.model_config = model.model_config self.model = model + # The executor drives the load balancer's iterations and teardown; this + # harness does it in `create_run_pack` and `shutdown_moe_load_balancer`. + self.moe_load_balancer = ( + load_balancer if isinstance(load_balancer, BaseMoeLoadBalancer) else None + ) self.layer_indices = layer_indices @staticmethod @@ -774,6 +786,11 @@ def create_run_pack( hc_mult = getattr(pretrained_config, "hc_mult", None) if hc_mult is not None: hidden_states = hidden_states.unsqueeze(1).expand(-1, hc_mult, -1).contiguous() + # Only layer 0 takes the plain residual. A slice starting mid-model + # enters with the resolved state that the previous layer returns. + first_layer = self.model.model.layers[self.layer_indices[0]] + if not getattr(first_layer, "is_first_layer", True): + hidden_states = HCState.resolved(hidden_states) kwargs["input_ids"] = torch.randint( 0, pretrained_config.vocab_size, @@ -801,7 +818,7 @@ def run_pack(*, check=False): discard_nccl_window_tensor_outputs((position_ids, hidden_states, residual)), ): get_model_extra_attrs()["attention_metadata"] = weakref.ref(attn_metadata) - with torch.inference_mode(): + with torch.inference_mode(), MoeLoadBalancerIterContext(self.moe_load_balancer): hidden_states_out, residual_out = self.model( position_ids, hidden_states, attn_metadata, residual, **kwargs ) @@ -816,6 +833,12 @@ def run_pack(*, check=False): return run_pack + def shutdown_moe_load_balancer(self) -> None: + """Release the MoE load balancer's resources; collective over the EP ranks.""" + if self.moe_load_balancer is not None: + self.moe_load_balancer.shutdown(local_safe=True) + self.moe_load_balancer = None + @contextlib.contextmanager def replace_routing_method_ctx(self, balance_method: BalanceMethod, balance_ratio: float): if balance_method == BalanceMethod.NotModified: From 3e698214249b13d905439ed5821b39962bc46c6f Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 05:11:46 -0700 Subject: [PATCH 12/59] [None][perf] Fuse HALO-Q and the in-switch TMA copy into one kernel Per-layer dynamic EPLB used to launch HALO-Q and the GPU-direct TMA copy of its plan as two kernels on the copy stream, with an event between them that MAIN waited on for the routes. Both now run in one kernel of auxiliary_sms CTAs (one per SM): - Every CTA first takes its HALO-Q role (CTA 0 plans, the others write routes), then adds one to a per-layer route epoch with a system-scope release, waits until all CTAs have done so, and copies the plan. - MAIN orders its first route read after the routes with cuStreamWaitValue64 on that epoch instead of the plan-ready event, so the wait holds no SM. - The kernel launches in the remote memory-sync domain, so fences of default-domain kernels do not wait for the copy's NVLink traffic. - Binding the copy loads the kernel and reserves its local memory; doing that at the first launch would wait for the device and could deadlock against kernels that spin on later work. The HALO-Q kernel body moves to moeRebalanceHaloQKernel.cuh as a device function shared by the standalone scheduler kernel (still used for the startup warmup) and the fused kernel. The standalone copy submit op is replaced by moe_rebalance_tma_submit_fused. auxiliary_sms now sizes the fused grid and the copy grid alike, so it must be between 1 and 8. Signed-off-by: jintaop --- .../dynamicEplb/moeRebalanceHaloQ.cu | 1949 +--------------- .../dynamicEplb/moeRebalanceHaloQ.h | 3 + .../dynamicEplb/moeRebalanceHaloQKernel.cuh | 1994 +++++++++++++++++ .../dynamicEplb/moeRebalancePlaceholder.cpp | 8 +- .../dynamicEplb/moeRebalanceTma.cu | 235 +- .../loadBalance/dynamicEplb/moeRebalanceTma.h | 18 +- .../moe/loadBalance/moeRebalanceHaloQOp.cpp | 38 +- .../moe/loadBalance/moeRebalanceHaloQOp.h | 42 + .../moe/loadBalance/moeRebalanceTmaOp.cpp | 57 +- examples/wide_ep/ep_load_balancer/README.md | 4 +- .../moe/fused_moe/MOE_DEVELOPER_GUIDE.md | 8 +- .../mega_moe/rebalance_slot_scheduler.py | 82 +- .../cuda_scheduler/runtime.py | 102 +- .../megamoe_scheduler/sami/hierarchical.py | 39 +- tensorrt_llm/llmapi/llm_args.py | 6 +- .../schemas/trtllm-serve-config.schema.json | 6 +- .../trtllm-serve-disagg-config.schema.json | 6 +- .../_torch/moe/test_dynamic_eplb_contracts.py | 112 +- tests/unittest/llmapi/test_llm_args.py | 19 +- .../test_trtllm_serve_config_schemas.py | 6 +- 20 files changed, 2599 insertions(+), 2135 deletions(-) create mode 100644 cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh create mode 100644 cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.h diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.cu b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.cu index 66d096e79813..f48c438602c6 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.cu +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.cu @@ -17,10 +17,8 @@ #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h" -#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlanPolicy.h" +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh" -#include -#include #include #include @@ -36,1746 +34,6 @@ namespace kernels namespace { -constexpr int kMaxEp = 32; -constexpr int kMaxExperts = 384; -constexpr int kMaxBroadcasts = kMaxExperts; -constexpr int kHelperMaskWords = (kMaxBroadcasts + 31) / 32; -constexpr int kMinGroup = 4; -constexpr int kSymPayloadOffset = 64; -using Epoch = std::uint32_t; - -enum PlanField -{ - kPlanExpert = 0, - kPlanLevel = 1, - kPlanGroupBegin = 2, - kPlanGroupSize = 3, - kPlanOwner = 4, - kPlanHelper = 5, -}; - -// Adapt placement storage to the copy layer without exposing scheduler -// internals. -__device__ __forceinline__ sami_copy::PlacementView copy_placement_view(int const* plan) -{ - return {plan[1], kMaxBroadcasts, plan + 2 + kPlanExpert * kMaxBroadcasts, plan + 2 + kPlanLevel * kMaxBroadcasts, - plan + 2 + kPlanGroupBegin * kMaxBroadcasts, plan + 2 + kPlanGroupSize * kMaxBroadcasts, - plan + 2 + kPlanOwner * kMaxBroadcasts, plan + 2 + kPlanHelper * kMaxBroadcasts}; -} - -__device__ __forceinline__ int imin(int a, int b) -{ - return a < b ? a : b; -} - -__device__ __forceinline__ int imax(int a, int b) -{ - return a > b ? a : b; -} - -__device__ __forceinline__ Epoch fetch_add_acq_rel_device(int* address, Epoch value) -{ - auto* valueAddress = reinterpret_cast(address); - cuda::atomic_ref atom(*valueAddress); - return atom.fetch_add(value, cuda::memory_order_acq_rel); -} - -__device__ __forceinline__ Epoch load_acquire_device(int* address) -{ - auto* valueAddress = reinterpret_cast(address); - cuda::atomic_ref atom(*valueAddress); - return atom.load(cuda::memory_order_acquire); -} - -__device__ __forceinline__ Epoch load_relaxed_device(int* address) -{ - auto* valueAddress = reinterpret_cast(address); - cuda::atomic_ref atom(*valueAddress); - return atom.load(cuda::memory_order_relaxed); -} - -__device__ __forceinline__ void store_relaxed_device(int* address, Epoch value) -{ - auto* valueAddress = reinterpret_cast(address); - cuda::atomic_ref atom(*valueAddress); - atom.store(value, cuda::memory_order_relaxed); -} - -__device__ __forceinline__ void store_release_device(int* address, Epoch value) -{ - auto* valueAddress = reinterpret_cast(address); - cuda::atomic_ref atom(*valueAddress); - atom.store(value, cuda::memory_order_release); -} - -__device__ __forceinline__ Epoch load_acquire_system(int* address) -{ - auto* valueAddress = reinterpret_cast(address); - cuda::atomic_ref atom(*valueAddress); - return atom.load(cuda::memory_order_acquire); -} - -__device__ __forceinline__ Epoch load_relaxed_system(int* address) -{ - auto* valueAddress = reinterpret_cast(address); - cuda::atomic_ref atom(*valueAddress); - return atom.load(cuda::memory_order_relaxed); -} - -__device__ __forceinline__ void store_release_system(int* address, Epoch value) -{ - auto* valueAddress = reinterpret_cast(address); - cuda::atomic_ref atom(*valueAddress); - atom.store(value, cuda::memory_order_release); -} - -__device__ __forceinline__ bool epoch_before(Epoch observed, Epoch expected) -{ - // RFC-1982-style serial comparison: all arithmetic remains unsigned and - // is unambiguous while peers differ by fewer than half the epoch space. - return observed != expected && expected - observed < (Epoch{1} << 31); -} - -__device__ __forceinline__ void record_status(int* address, Epoch epoch) -{ - // Zero is reserved for "never failed". Preserve a nonzero sticky error at - // the one modular epoch whose bit pattern is zero. - Epoch const value = epoch == 0 ? ~Epoch{0} : epoch; - atomicExch(reinterpret_cast(address), value); -} - -__device__ __forceinline__ int* plan_field(int* plan, int field) -{ - return plan + 2 + field * kMaxBroadcasts; -} - -__device__ __forceinline__ bool group_contains(int rank, int begin, int size) -{ - return rank >= begin && rank < begin + size; -} - -__host__ __device__ __forceinline__ unsigned rank_group_mask(int begin, int size) -{ - return size == 32 ? 0xffffffffu : ((1u << size) - 1u) << begin; -} - -__host__ __device__ __forceinline__ int hierarchy_depth(int ep) -{ - int depth = 1; - for (int level = 1; (1 << level) <= ep; ++level) - { - int groups = 1 << level; - if (ep % groups || ep / groups < kMinGroup) - break; - ++depth; - } - return depth; -} - -/* Distinct experts and total receive capacity independently bound broadcasts. - * Actual S is retained in the public outputs; only scratch uses this bound. */ -__host__ __device__ __forceinline__ int broadcast_capacity(int ep, int experts, int helpers) -{ - int const bounded_helpers = helpers < experts ? helpers : experts; - /* Every hierarchy level divides ep exactly. This cancels two integer - * divisions without changing capacity, including non-power-of-two EP. */ - int const groups = 1 << (hierarchy_depth(ep) - 1); - int const receive_bound = groups * bounded_helpers; - return receive_bound < experts ? receive_bound : experts; -} - -__device__ bool grid_rendezvous( - int* sync, int blocks, unsigned long long spin_cycles, int* status, Epoch* completedEpoch) -{ - __shared__ Epoch generation; - __shared__ int success; - /* The block barrier orders every writer's stores before thread zero's - * device-scope release advertises this CTA. */ - __syncthreads(); - if (threadIdx.x == 0) - { - success = 1; - Epoch const current = load_relaxed_device(sync + 1); - Epoch const target = current + 1; - generation = target; - Epoch const old = fetch_add_acq_rel_device(sync, 1); - if (old == static_cast(blocks - 1)) - { - // The last CTA resets arrivals before publishing the next modular - // epoch. Stream ordering prevents a new launch from entering until - // every CTA in this launch has observed the release. - store_relaxed_device(sync, 0); - store_release_device(sync + 1, target); - } - else if (old >= static_cast(blocks)) - { - record_status(status + 1, target); - success = 0; - } - else - { - unsigned long long const start = clock64(); - Epoch observed = load_relaxed_device(sync + 1); - while (epoch_before(observed, target)) - { - if (clock64() - start > spin_cycles) - { - record_status(status + 1, target); - success = 0; - break; - } - observed = load_relaxed_device(sync + 1); - } - if (!epoch_before(observed, target)) - (void) load_acquire_device(sync + 1); - } - } - __syncthreads(); - *completedEpoch = generation; - return success != 0; -} - -__device__ void local_histogram_pass(int const* routes, int* partial, int route_count, int expert_count, int ctas) -{ - extern __shared__ int shared[]; - int const tid = threadIdx.x; - int const lane = tid & 31; - int const warp = tid >> 5; - int const warps = blockDim.x >> 5; - int const bins = expert_count + 1; - int const shards = warps > 8 ? 8 : warps; - int const shared_count = shards * bins; - for (int index = tid; index < shared_count; index += blockDim.x) - shared[index] = 0; - __syncthreads(); - - int const cta_span = (route_count + ctas - 1) / ctas; - int const cta_begin = blockIdx.x * cta_span; - int const cta_end = imin(route_count, cta_begin + cta_span); - int const warp_span = ((cta_span + warps * 32 - 1) / (warps * 32)) * 32; - int const begin = cta_begin + warp * warp_span; - int const end = imin(cta_end, begin + warp_span); - /* Issue a small load group before consuming it to expose memory-level - * parallelism. The binning is unchanged, and atomicAdd regrouping is - * safe because integer addition commutes. */ - constexpr int kBatch = 4; - for (int base_index = begin; base_index < end; base_index += 32 * kBatch) - { - int expert[kBatch]; -#pragma unroll - for (int step = 0; step < kBatch; ++step) - { - int const index = base_index + step * 32 + lane; - expert[step] = index < end ? routes[index] : -1; - } -#pragma unroll - for (int step = 0; step < kBatch; ++step) - { - int const index = base_index + step * 32 + lane; - if (index < end) - { - int const value = expert[step]; - int const bin = (value >= 0 && value < expert_count) ? value : expert_count; - atomicAdd(shared + (warp & 7) * bins + bin, 1); - } - } - } - __syncthreads(); - - for (int expert = tid; expert < expert_count; expert += blockDim.x) - { - int count = 0; - for (int shard = 0; shard < shards; ++shard) - count += shared[shard * bins + expert]; - partial[blockIdx.x * expert_count + expert] = count; - } -} - -/* Re-rank all local routes on the C-1 worker CTAs while CTA0 performs the - * cross-rank exchange and planner. route_aux[0:2] is a reusable worker - * rendezvous; the following NWORK*(E+1) entries hold per-worker totals. - * Per-warp cursors remain in this CTA's shared memory until materialization. */ -__device__ bool stable_worker_ordinal_pass(int const* routes, int* ordinal, int* route_aux, int* status, - int route_count, int expert_count, int worker, int workers, unsigned long long spin_cycles) -{ - extern __shared__ int shared[]; - int const tid = threadIdx.x; - int const lane = tid & 31; - int const warp = tid >> 5; - int const warps = blockDim.x >> 5; - int const bins = expert_count + 1; - int const shared_count = warps * bins; - for (int index = tid; index < shared_count; index += blockDim.x) - shared[index] = 0; - __syncthreads(); - - int const worker_span = (route_count + workers - 1) / workers; - int const worker_begin = worker * worker_span; - int const worker_end = imin(route_count, worker_begin + worker_span); - int const warp_span = ((worker_span + warps * 32 - 1) / (warps * 32)) * 32; - int const begin = worker_begin + warp * warp_span; - int const end = imin(worker_end, begin + warp_span); - for (int base_index = begin; base_index < end; base_index += 32) - { - int const index = base_index + lane; - bool const active = index < end; - int const expert = active ? routes[index] : -1; - // Padding has its own key so it cannot contribute to real/invalid bins. - int const bin = !active ? bins : ((expert >= 0 && expert < expert_count) ? expert : expert_count); - unsigned const peers = __match_any_sync(0xffffffffu, bin); - int const leader = __ffs(peers) - 1; - int const rank = __popc(peers & ((1u << lane) - 1u)); - int base = 0; - if (active && lane == leader) - base = atomicAdd(shared + warp * bins + bin, __popc(peers)); - // One collective for the entire warp; each lane selects its own leader. - base = __shfl_sync(0xffffffffu, base, leader); - if (active) - ordinal[index] = base + rank; - } - __syncthreads(); - - int* worker_totals = route_aux + 2; - for (int bin = tid; bin < bins; bin += blockDim.x) - { - int total = 0; - for (int local_warp = 0; local_warp < warps; ++local_warp) - total += shared[local_warp * bins + bin]; - worker_totals[worker * bins + bin] = total; - } - - /* The first all-CTA rendezvous guarantees every worker is resident, so a - * worker-only global barrier cannot wait on an unscheduled CTA. */ - Epoch workerEpoch; - if (!grid_rendezvous(route_aux, workers, spin_cycles, status, &workerEpoch)) - return false; - - for (int bin = tid; bin < bins; bin += blockDim.x) - { - int cursor = 0; - for (int previous = 0; previous < worker; ++previous) - cursor += worker_totals[previous * bins + bin]; - for (int local_warp = 0; local_warp < warps; ++local_warp) - { - int index = local_warp * bins + bin; - int count = shared[index]; - shared[index] = cursor; - cursor += count; - } - } - __syncthreads(); - return true; -} - -template -__device__ bool exchange_counts(unsigned long long const* peer_bases, int const* partial, int* shared, int ep, - int expert_count, int local_rank, int ctas, Epoch epoch, unsigned long long spin_cycles, int* status) -{ - int const exchange_ep = FixedEp > 0 ? FixedEp : ep; - int const tid = threadIdx.x; - __shared__ int wait_failed; - if (tid == 0) - wait_failed = 0; - /* Keep the simple accumulation loop; the compiler provides the required - * pipelining without explicit batching. */ - for (int expert = tid; expert < expert_count; expert += blockDim.x) - { - int count = 0; - for (int cta = 0; cta < ctas; ++cta) - count += partial[cta * expert_count + expert]; - shared[expert] = count; - } - __syncthreads(); - - int const parity = epoch & 1; - /* One warp owns each destination row. Its warp barrier orders every - * lane's stores before lane zero's system-scope release publishes the - * complete row to the peer acquire. */ - int const warp = tid >> 5; - int const lane = tid & 31; - for (int destination = warp; destination < exchange_ep; destination += blockDim.x >> 5) - { - int* peer = reinterpret_cast(peer_bases[destination]); - int* row = peer + kSymPayloadOffset + (parity * exchange_ep + local_rank) * expert_count; - /* Keep scalar remote stores for the publication path; the local mirror - * copy below uses vector transfers when alignment permits. */ - for (int expert = lane; expert < expert_count; expert += 32) - row[expert] = shared[expert]; - /* Warp synchronization orders every lane's row stores before lane - * zero resumes; its system-scope release then publishes that complete - * transitive happens-before set to the peer acquire. */ - __syncwarp(); - if (lane == 0) - store_release_system(peer + local_rank, epoch); - } - __syncthreads(); - - int* local = reinterpret_cast(peer_bases[local_rank]); - if (tid < exchange_ep) - { - unsigned long long start = clock64(); - Epoch observed = load_relaxed_system(local + tid); - while (epoch_before(observed, epoch)) - { - if (clock64() - start > spin_cycles) - { - record_status(status + 0, epoch); - atomicExch(&wait_failed, 1); - break; - } - observed = load_relaxed_system(local + tid); - } - if (!epoch_before(observed, epoch)) - (void) load_acquire_system(local + tid); - } - __syncthreads(); - if (wait_failed) - return false; - int const* payload = local + kSymPayloadOffset + parity * exchange_ep * expert_count; - int const total = exchange_ep * expert_count; - /* Vectorize the local mirror copy when shape and alignment permit. */ - if ((total & 3) == 0 && (reinterpret_cast(payload) & 15) == 0) - { - int4 const* source = reinterpret_cast(payload); - int4* destination = reinterpret_cast(shared); - int const vectors = total >> 2; - for (int vector = tid; vector < vectors; vector += blockDim.x) - destination[vector] = source[vector]; - } - else - { - for (int index = tid; index < total; index += blockDim.x) - shared[index] = payload[index]; - } - __syncthreads(); - return true; -} - -__device__ void prepare_halo_m_sparse8(int const* counts, int* histogram, int* load, int* candidates, int experts) -{ - int const tid = threadIdx.x; - int const warp = tid >> 5; - int const lane = tid & 31; - if (warp >= 8) - return; - int const home = experts / 8; - int local_index[2]; - int local_mass[2]; - int rank_load = 0; -#pragma unroll - for (int round = 0; round < 2; ++round) - { - int local = lane + 32 * round; - int mass = -1; - if (local < home) - { - mass = 0; - int expert = warp * home + local; -#pragma unroll - for (int source = 0; source < 8; ++source) - mass += counts[source * experts + expert]; - rank_load += mass; - } - local_index[round] = local; - local_mass[round] = mass; - } - rank_load = __reduce_add_sync(0xffffffffu, rank_load); - if (lane == 0) - load[warp] = rank_load; - for (int slot = 0; slot < 4; ++slot) - { - int best_mass = -1; - int best_local = INT_MAX; -#pragma unroll - for (int round = 0; round < 2; ++round) - { - int mass = local_mass[round]; - int local = local_index[round]; - if (mass > best_mass || (mass == best_mass && local < best_local)) - { - best_mass = mass; - best_local = local; - } - } - int warp_mass = __reduce_max_sync(0xffffffffu, best_mass); - int winner = __reduce_min_sync(0xffffffffu, best_mass == warp_mass ? best_local : INT_MAX); - if (lane == 0) - { - int expert = warp * home + winner; - candidates[warp * 4 + slot] = expert; - histogram[expert] = warp_mass; - } -#pragma unroll - for (int round = 0; round < 2; ++round) - if (local_index[round] == winner) - local_mass[round] = -1; - } - asm volatile("bar.sync 4, 256;" ::: "memory"); -} - -/* Build HALO-M's dense histogram, rank loads and per-owner hot lists with one - * warp per rank. The former thread-0 implementation dominated the kernel. */ -template -__device__ void prepare_halo_m(int const* counts, int* histogram, int* load, int* candidates, int* recv_used, - int* send_used, int* next_index, int* frozen, int ep, int experts, int helpers) -{ - int const prepare_ep = FixedEp > 0 ? FixedEp : ep; - int const prepare_helpers = imin(FixedHelpers > 0 ? FixedHelpers : helpers, experts / prepare_ep); - int const tid = threadIdx.x; - int const warp = tid >> 5; - int const lane = tid & 31; - int const home = experts / prepare_ep; - for (int expert = tid; expert < experts; expert += blockDim.x) - { - int mass = 0; -#pragma unroll - for (int rank = 0; rank < prepare_ep; ++rank) - mass += counts[rank * experts + expert]; - histogram[expert] = mass; - } - __syncthreads(); - - for (int owner = warp; owner < prepare_ep; owner += blockDim.x >> 5) - { - int rank_load = 0; - for (int local = lane; local < home; local += 32) - rank_load += histogram[owner * home + local]; - rank_load = __reduce_add_sync(0xffffffffu, rank_load); - if (lane == 0) - { - load[owner] = rank_load; - recv_used[owner] = 0; - send_used[owner] = 0; - next_index[owner] = 0; - frozen[owner] = 0; - } - - /* EP2 is the largest home: 192 experts, six entries per lane. */ - constexpr int rounds = FixedEp == 8 ? 2 : 6; - int local_index[rounds]; - int local_mass[rounds]; -#pragma unroll - for (int round = 0; round < rounds; ++round) - { - int local = lane + 32 * round; - local_index[round] = local; - local_mass[round] = local < home ? histogram[owner * home + local] : -1; - } - for (int slot = 0; slot < prepare_helpers; ++slot) - { - int best_mass = -1; - int best_local = INT_MAX; -#pragma unroll - for (int round = 0; round < rounds; ++round) - { - int mass = local_mass[round]; - int local = local_index[round]; - if (mass > best_mass || (mass == best_mass && local < best_local)) - { - best_mass = mass; - best_local = local; - } - } - int warp_mass = __reduce_max_sync(0xffffffffu, best_mass); - int winner = __reduce_min_sync(0xffffffffu, best_mass == warp_mass ? best_local : INT_MAX); - if (lane == 0) - candidates[owner * prepare_helpers + slot] = owner * home + winner; -#pragma unroll - for (int round = 0; round < rounds; ++round) - if (local_index[round] == winner) - local_mass[round] = -1; - } - } - __syncthreads(); -} - -/* EP<=8 fast path: every participating lane owns one rank's complete - * mutable HALO-M state. Group and destination decisions use only shuffles, - * avoiding repeated shared-memory rank-state traffic and warp barriers. */ -template -__device__ __forceinline__ int plan_halo_m_register8_impl( - int const* histogram, int lane_load, int total, int* quota, int const* candidates, int* plan, int ep, int helpers) -{ - int const plan_ep = FixedEp > 0 ? FixedEp : ep; - int const plan_helpers = FixedHelpers > 0 ? FixedHelpers : helpers; - int const lane = threadIdx.x & 31; - constexpr unsigned mask = 0xffu; - int lane_recv = 0; - int lane_send = 0; - int lane_next = 0; - int lane_frozen = 0; - - int const target = total / plan_ep + (total % plan_ep != 0); - int const depth = hierarchy_depth(plan_ep); - int* p_expert = plan_field(plan, kPlanExpert); - int* p_level = plan_field(plan, kPlanLevel); - int* p_begin = plan_field(plan, kPlanGroupBegin); - int* p_size = plan_field(plan, kPlanGroupSize); - int* p_owner = plan_field(plan, kPlanOwner); - int* p_helper = plan_field(plan, kPlanHelper); - int broadcasts = 0; - - int trips = 0; - for (int guard = 0; guard < plan_ep * (plan_helpers + 2); ++guard) - { - ++trips; - bool const eligible = lane < plan_ep && !lane_frozen && lane_load > target; - int owner, owner_load; - if constexpr (NarrowKeys) - { - unsigned const key = __reduce_max_sync( - mask, eligible ? ((static_cast(lane_load) << 4) | (15u - static_cast(lane))) : 0u); - if (!key) - break; - owner = 15 - static_cast(key & 15u); - owner_load = static_cast(key >> 4); - } - else - { - owner_load = __reduce_max_sync(mask, eligible ? lane_load : -1); - if (owner_load < 0) - break; - owner = __reduce_min_sync(mask, eligible && lane_load == owner_load ? lane : INT_MAX); - } - - unsigned const packed - = __shfl_sync(mask, (static_cast(lane_send) << 4) | static_cast(lane_next), owner, 8); - int const owner_send = static_cast(packed >> 4); - int const owner_next = static_cast(packed & 15u); - if (owner_send >= plan_helpers || owner_next >= plan_helpers) - { - if (lane == owner) - lane_frozen = 1; - continue; - } - - /* owner and owner_next are now warp-uniform, so every lane can read - * candidates[] and histogram[] directly. Both are written by - * prepare_halo_m_sparse8 before its bar.sync and never inside this - * loop, so a uniform load returns exactly what the owner would have - * broadcast. The increment stays owner-only and stays ahead of the - * mass test, as before. */ - int const expert = candidates[owner * plan_helpers + owner_next]; - if (lane == owner) - ++lane_next; - int const mass = histogram[expert]; - if (mass == 0) - { - if (lane == owner) - lane_frozen = 1; - continue; - } - - int give = imin(owner_load - target, mass); - if (give <= 0) - { - if (lane == owner) - lane_frozen = 1; - continue; - } - - int positive_target = imax(target, 1); - int need = imax(1, mass / positive_target + (mass % positive_target != 0)); - int wanted = depth - 1; - for (int level = depth - 1; level >= 0; --level) - { - if (plan_ep / (1 << level) >= need) - { - wanted = level; - break; - } - } - - int chosen_level = -1; - int chosen_begin = -1; - int chosen_size = 0; - if constexpr (FixedEp == 8 && FixedHelpers == 4) - { - int slack = lane != owner ? imax(0, target - lane_load) : 0; - /* Two independent full-mask sums instead of one lane-varying-mask - * reduce followed by two broadcasts. Integer addition is exact, so - * masking the addend gives the same half sums; the two reduces do - * not depend on each other, so the chain is depth 1 instead of 3, - * and the redux.sync no longer runs with a per-lane membermask. */ - int room0 = __reduce_add_sync(mask, lane < 4 ? slack : 0); - int room4 = __reduce_add_sync(mask, lane >= 4 ? slack : 0); - unsigned receive_mask = __ballot_sync(mask, lane_recv < plan_helpers); - bool valid0 = (receive_mask & 0x0fu) == 0x0fu && room0 > 0; - bool valid4 = (receive_mask & 0xf0u) == 0xf0u && room4 > 0; - bool want_half = mass <= 4 * target; - if (want_half && (valid0 || valid4)) - { - chosen_level = 1; - chosen_begin = valid0 && (!valid4 || room0 >= room4) ? 0 : 4; - chosen_size = 4; - } - else if (receive_mask == mask && room0 + room4 > 0) - { - chosen_level = 0; - chosen_begin = 0; - chosen_size = 8; - } - } - else - { - for (int level = wanted; level >= 0 && chosen_level < 0; --level) - { - int size = plan_ep / (1 << level); - bool active = lane < plan_ep; - int begin = (lane / size) * size; - unsigned group_mask = rank_group_mask(begin, size); - int local_room = active && lane != owner ? imax(0, target - lane_load) : 0; - int room = 0; - if (active) - room = __reduce_add_sync(group_mask, local_room); - unsigned receive_mask = __ballot_sync(mask, active && lane_recv < plan_helpers); - bool group_leader = active && lane == begin && (receive_mask & group_mask) == group_mask; - unsigned leader_mask = __ballot_sync(mask, group_leader); - int best_room = -1; - int best_begin = INT_MAX; - if (leader_mask) - { - int const first_leader = __ffs(leader_mask) - 1; - if constexpr (NarrowKeys) - { - unsigned key = 0; - if (group_leader) - key = __reduce_max_sync( - leader_mask, (static_cast(room) << 4) | (15u - static_cast(begin))); - key = __shfl_sync(mask, key, first_leader, 8); - best_room = static_cast(key >> 4); - best_begin = 15 - static_cast(key & 15u); - } - else - { - if (group_leader) - { - best_room = __reduce_max_sync(leader_mask, room); - best_begin = __reduce_min_sync(leader_mask, room == best_room ? begin : INT_MAX); - } - best_room = __shfl_sync(mask, best_room, first_leader, 8); - best_begin = __shfl_sync(mask, best_begin, first_leader, 8); - } - } - if (best_begin != INT_MAX && best_room > 0) - { - chosen_level = level; - chosen_begin = best_begin; - chosen_size = size; - } - } - } - if (chosen_level < 0) - { - if (lane == owner) - lane_frozen = 1; - continue; - } - if (broadcasts >= kMaxBroadcasts) - break; - - int lane_quota = lane == owner ? mass : 0; - int remaining = give; - unsigned visited = 0; - for (int step = 0; step < chosen_size && remaining > 0; ++step) - { - bool can_receive = lane >= chosen_begin && lane < chosen_begin + chosen_size && lane != owner - && !(visited & (1u << lane)); - int destination, destination_load; - if constexpr (NarrowKeys) - { - unsigned const key = __reduce_min_sync(mask, - can_receive ? ((static_cast(lane_load) << 4) | static_cast(lane)) - : 0xffffffffu); - if (key == 0xffffffffu) - break; - destination = static_cast(key & 15u); - destination_load = static_cast(key >> 4); - } - else - { - destination_load = __reduce_min_sync(mask, can_receive ? lane_load : INT_MAX); - destination = __reduce_min_sync(mask, can_receive && lane_load == destination_load ? lane : INT_MAX); - if (destination == INT_MAX) - break; - } - visited |= 1u << destination; - int room = target - destination_load; - if (room <= 0) - break; - int amount = imin(room, remaining); - if (lane == owner) - lane_quota -= amount; - if (lane == destination) - lane_quota += amount; - remaining -= amount; - } - - int moved = give - remaining; - if (moved <= 0) - { - if (lane == owner) - lane_frozen = 1; - continue; - } - - int b = broadcasts++; - if (lane == 0) - { - p_expert[b] = expert; - p_level[b] = chosen_level; - p_begin[b] = chosen_begin; - p_size[b] = chosen_size; - p_owner[b] = owner; - p_helper[b] = -1; - } - if (lane < plan_ep) - quota[b * plan_ep + lane] = lane_quota; - - bool in_group = lane >= chosen_begin && lane < chosen_begin + chosen_size; - if (in_group) - ++lane_recv; - if (lane == owner) - { - ++lane_send; - lane_load -= moved; - } - else if (in_group) - { - lane_load += lane_quota; - } - } - return broadcasts; -} - -/* Select the key representation once, from the actual complete rank load. - * This guard keeps owner/destination loads and group-room keys representable. - * Keep the wide path for large valid int32 inputs without putting a runtime - * representation branch on each iteration of either planning loop. */ -template -__device__ __forceinline__ int plan_halo_m_register8( - int const* histogram, int const* load, int* quota, int const* candidates, int* plan, int ep, int helpers) -{ - int const plan_ep = FixedEp > 0 ? FixedEp : ep; - int const lane = threadIdx.x & 31; - int const lane_load = lane < plan_ep ? load[lane] : 0; - int const total = __reduce_add_sync(0xffu, lane_load); - if (total <= 0x0fffffff) - return plan_halo_m_register8_impl( - histogram, lane_load, total, quota, candidates, plan, ep, helpers); - return plan_halo_m_register8_impl( - histogram, lane_load, total, quota, candidates, plan, ep, helpers); -} - -/* Exact HALO-M state machine executed by warp 0. Each lane owns one rank; - * warp reductions replace all serial rank/group/destination scans. */ -__device__ int plan_halo_m_warp(int const* histogram, int* load, int* quota, int const* candidates, int* recv_used, - int* send_used, int* next_index, int* frozen, int* plan, int ep, int experts, int helpers) -{ - int const candidate_count = imin(helpers, experts / ep); - int const lane = threadIdx.x & 31; - int const reduce_width = ep <= 8 ? 8 : 32; - int const reduce_start = reduce_width >> 1; - unsigned const reduce_mask = ep <= 8 ? 0xffu : 0xffffffffu; - int total = lane < ep ? load[lane] : 0; - for (int offset = reduce_start; offset; offset >>= 1) - { - int const other = __shfl_down_sync(reduce_mask, total, offset, reduce_width); - if (lane + offset < reduce_width) - total += other; - } - total = __shfl_sync(reduce_mask, total, 0, reduce_width); - int const target = total / ep + (total % ep != 0); - int const depth = hierarchy_depth(ep); - int* p_expert = plan_field(plan, kPlanExpert); - int* p_level = plan_field(plan, kPlanLevel); - int* p_begin = plan_field(plan, kPlanGroupBegin); - int* p_size = plan_field(plan, kPlanGroupSize); - int* p_owner = plan_field(plan, kPlanOwner); - int* p_helper = plan_field(plan, kPlanHelper); - int broadcasts = 0; - - for (int guard = 0; guard < ep * (candidate_count + 2); ++guard) - { - bool eligible = lane < ep && !frozen[lane] && load[lane] > target; - int best_load = eligible ? load[lane] : -1; - int best_rank = eligible ? lane : INT_MAX; - for (int offset = reduce_start; offset; offset >>= 1) - { - int other_load = __shfl_down_sync(reduce_mask, best_load, offset, reduce_width); - int other_rank = __shfl_down_sync(reduce_mask, best_rank, offset, reduce_width); - if (lane + offset < reduce_width - && (other_load > best_load || (other_load == best_load && other_rank < best_rank))) - { - best_load = other_load; - best_rank = other_rank; - } - } - int owner = __shfl_sync(reduce_mask, best_rank, 0, reduce_width); - if (owner == INT_MAX) - break; - /* One lane owns the capacity check and shared cursor transition. - * Publish either the chosen expert or -1 for exhaustion explicitly; - * no reader can observe a partially advanced cursor. */ - int expert = -1; - if (lane == 0) - { - if (send_used[owner] >= helpers || next_index[owner] >= candidate_count) - { - frozen[owner] = 1; - } - else - { - int const candidate_index = next_index[owner]++; - expert = candidates[owner * candidate_count + candidate_index]; - } - } - expert = __shfl_sync(reduce_mask, expert, 0, reduce_width); - __syncwarp(reduce_mask); - if (expert < 0) - continue; - int mass = histogram[expert]; - if (mass == 0) - { - if (lane == 0) - frozen[owner] = 1; - __syncwarp(reduce_mask); - continue; - } - int give = imin(load[owner] - target, mass); - if (give <= 0) - { - if (lane == 0) - frozen[owner] = 1; - __syncwarp(reduce_mask); - continue; - } - - int need = imax(1, mass / imax(target, 1) + (mass % imax(target, 1) != 0)); - int wanted = depth - 1; - for (int level = depth - 1; level >= 0; --level) - { - if (ep / (1 << level) >= need) - { - wanted = level; - break; - } - } - int chosen_level = -1; - int chosen_begin = -1; - int chosen_size = 0; - for (int level = wanted; level >= 0 && chosen_level < 0; --level) - { - int size = ep / (1 << level); - int begin = lane * size; - bool active = begin < ep; - bool affordable = active; - int room = 0; - if (active) - { - for (int rank = begin; rank < begin + size; ++rank) - { - affordable &= recv_used[rank] < helpers; - if (rank != owner) - room += imax(0, target - load[rank]); - } - } - int candidate_room = affordable ? room : -1; - int candidate_begin = affordable ? begin : INT_MAX; - for (int offset = reduce_start; offset; offset >>= 1) - { - int other_room = __shfl_down_sync(reduce_mask, candidate_room, offset, reduce_width); - int other_begin = __shfl_down_sync(reduce_mask, candidate_begin, offset, reduce_width); - if (lane + offset < reduce_width - && (other_room > candidate_room || (other_room == candidate_room && other_begin < candidate_begin))) - { - candidate_room = other_room; - candidate_begin = other_begin; - } - } - int best_room = __shfl_sync(reduce_mask, candidate_room, 0, reduce_width); - int best_begin = __shfl_sync(reduce_mask, candidate_begin, 0, reduce_width); - if (best_begin != INT_MAX && best_room > 0) - { - chosen_level = level; - chosen_begin = best_begin; - chosen_size = size; - } - } - if (chosen_level < 0) - { - if (lane == 0) - frozen[owner] = 1; - __syncwarp(reduce_mask); - continue; - } - if (broadcasts >= kMaxBroadcasts) - break; - - int b = broadcasts++; - if (lane == 0) - { - p_expert[b] = expert; - p_level[b] = chosen_level; - p_begin[b] = chosen_begin; - p_size[b] = chosen_size; - p_owner[b] = owner; - p_helper[b] = -1; - } - if (lane < ep) - quota[b * ep + lane] = 0; - __syncwarp(reduce_mask); - if (lane == 0) - quota[b * ep + owner] = mass; - __syncwarp(reduce_mask); - - int remaining = give; - unsigned visited = 0; - for (int step = 0; step < chosen_size && remaining > 0; ++step) - { - bool can_receive = lane >= chosen_begin && lane < chosen_begin + chosen_size && lane != owner - && !(visited & (1u << lane)); - int cold_load = can_receive ? load[lane] : INT_MAX; - int cold_rank = can_receive ? lane : INT_MAX; - for (int offset = reduce_start; offset; offset >>= 1) - { - int other_load = __shfl_down_sync(reduce_mask, cold_load, offset, reduce_width); - int other_rank = __shfl_down_sync(reduce_mask, cold_rank, offset, reduce_width); - if (lane + offset < reduce_width - && (other_load < cold_load || (other_load == cold_load && other_rank < cold_rank))) - { - cold_load = other_load; - cold_rank = other_rank; - } - } - int destination = __shfl_sync(reduce_mask, cold_rank, 0, reduce_width); - if (destination == INT_MAX) - break; - visited |= 1u << destination; - int room = target - load[destination]; - if (room <= 0) - continue; - int amount = imin(room, remaining); - if (lane == 0) - { - quota[b * ep + owner] -= amount; - quota[b * ep + destination] += amount; - } - remaining -= amount; - __syncwarp(reduce_mask); - } - int moved = give - remaining; - if (moved <= 0) - { - --broadcasts; - if (lane == 0) - frozen[owner] = 1; - __syncwarp(reduce_mask); - continue; - } - if (lane >= chosen_begin && lane < chosen_begin + chosen_size) - ++recv_used[lane]; - if (lane == owner) - { - ++send_used[owner]; - load[owner] -= moved; - } - else if (lane >= chosen_begin && lane < chosen_begin + chosen_size) - { - load[lane] += quota[b * ep + lane]; - } - __syncwarp(reduce_mask); - } - return broadcasts; -} - -__device__ __forceinline__ int warp_sum(int value) -{ - for (int offset = 16; offset; offset >>= 1) - { - int const other = __shfl_down_sync(0xffffffffu, value, offset); - if ((threadIdx.x & 31) + offset < 32) - value += other; - } - return __shfl_sync(0xffffffffu, value, 0); -} - -__device__ __forceinline__ int warp_exclusive_scan(int value) -{ - int sum = value; - int const lane = threadIdx.x & 31; - for (int offset = 1; offset < 32; offset <<= 1) - { - int other = __shfl_up_sync(0xffffffffu, sum, offset); - if (lane >= offset) - sum += other; - } - return sum - value; -} - -__device__ void halo_q_repair(int receiver, int donor, int broadcasts, int const* order, int const* host_mask, - int const* counts, int const* p_expert, int* quota, int ep, int experts) -{ - int const lane = threadIdx.x & 31; - if (broadcasts > 32) - { - int local_pull = 0, local_give = 0; - for (int base = 0; base < broadcasts; base += 32) - { - int const b = base + lane < broadcasts ? order[base + lane] : -1; - if (b >= 0 && (host_mask[b] & (1u << receiver)) && (host_mask[b] & (1u << donor))) - { - int const expert = p_expert[b]; - int const cr = counts[receiver * experts + expert]; - int const cd = counts[donor * experts + expert]; - int const qr = quota[b * ep + receiver]; - int const qd = quota[b * ep + donor]; - local_pull += imin(imax(0, cr - qr), imax(0, qd - cd)); - local_give += imax(0, qr - cr); - } - } - int const delta = imin(warp_sum(local_pull), warp_sum(local_give)); - int pull_before_tiles = 0, give_before_tiles = 0; - for (int base = 0; base < broadcasts; base += 32) - { - int const b = base + lane < broadcasts ? order[base + lane] : -1; - int pull = 0, give = 0; - if (b >= 0 && (host_mask[b] & (1u << receiver)) && (host_mask[b] & (1u << donor))) - { - int const expert = p_expert[b]; - int const cr = counts[receiver * experts + expert]; - int const cd = counts[donor * experts + expert]; - int const qr = quota[b * ep + receiver]; - int const qd = quota[b * ep + donor]; - pull = imin(imax(0, cr - qr), imax(0, qd - cd)); - give = imax(0, qr - cr); - } - int const pull_before = pull_before_tiles + warp_exclusive_scan(pull); - int const give_before = give_before_tiles + warp_exclusive_scan(give); - int const take = imin(pull, imax(0, delta - pull_before)); - int const release = imin(give, imax(0, delta - give_before)); - pull_before_tiles += warp_sum(pull); - give_before_tiles += warp_sum(give); - if (b >= 0) - { - quota[b * ep + receiver] += take - release; - quota[b * ep + donor] += release - take; - } - __syncwarp(); - } - return; - } - int b = lane < broadcasts ? order[lane] : -1; - int pull = 0; - int give = 0; - if (b >= 0 && (host_mask[b] & (1u << receiver)) && (host_mask[b] & (1u << donor))) - { - int expert = p_expert[b]; - int cr = counts[receiver * experts + expert]; - int cd = counts[donor * experts + expert]; - int qr = quota[b * ep + receiver]; - int qd = quota[b * ep + donor]; - pull = imin(imax(0, cr - qr), imax(0, qd - cd)); - give = imax(0, qr - cr); - } - int delta = imin(warp_sum(pull), warp_sum(give)); - int pull_before = warp_exclusive_scan(pull); - int give_before = warp_exclusive_scan(give); - int pull_take = imin(pull, imax(0, delta - pull_before)); - int give_take = imin(give, imax(0, delta - give_before)); - if (b >= 0) - { - quota[b * ep + receiver] += pull_take - give_take; - quota[b * ep + donor] += give_take - pull_take; - } - __syncwarp(); -} - -__device__ __forceinline__ int warp_sum8(int value) -{ - return __reduce_add_sync(0xffu, value); -} - -__device__ __forceinline__ int warp_exclusive_scan8(int value) -{ - int sum = value; - int const lane = threadIdx.x & 7; - for (int offset = 1; offset < 8; offset <<= 1) - { - int other = __shfl_up_sync(0xffu, sum, offset, 8); - if (lane >= offset) - sum += other; - } - return sum - value; -} - -__device__ __forceinline__ void halo_q_direction8( - bool common, int receiver_count, int donor_count, int& receiver_quota, int& donor_quota) -{ - int pull = common ? imin(imax(0, receiver_count - receiver_quota), imax(0, donor_quota - donor_count)) : 0; - int give = common ? imax(0, receiver_quota - receiver_count) : 0; - int pull_total = warp_sum8(pull); - int give_total = warp_sum8(give); - int delta = imin(pull_total, give_total); - if (!delta) - return; - int pull_before = warp_exclusive_scan8(pull); - int give_before = warp_exclusive_scan8(give); - int pull_take = imin(pull, imax(0, delta - pull_before)); - int give_take = imin(give, imax(0, delta - give_before)); - receiver_quota += pull_take - give_take; - donor_quota += give_take - pull_take; -} - -/* Fuse both ordered directions of one rank pair. Each broadcast lane loads - * its two count/quota entries once, executes left<-right then right<-left in - * registers, and writes the final pair once. - * - * `b`, `hosts` and `expert` are the lane's own broadcast identity and do not - * change between phases or rounds, so the caller hoists them. Reading them - * here put a four-deep dependent load chain (order[lane] -> b -> host_mask[b] - * and the GLOBAL p_expert[b] -> expert -> counts[...]) at the head of all - * fourteen phases. */ -__device__ __forceinline__ void halo_q_pair8( - int left, int right, int b, int hosts, int expert, int const* counts, int* quota, int ep, int experts) -{ - int const lane = threadIdx.x & 31; - if (lane >= 8) - return; - bool common = false; - int count_left = 0; - int count_right = 0; - int quota_left = 0; - int quota_right = 0; - if (b >= 0) - { - common = (hosts & (1u << left)) && (hosts & (1u << right)); - if (common) - { - count_left = counts[left * experts + expert]; - count_right = counts[right * experts + expert]; - quota_left = quota[b * ep + left]; - quota_right = quota[b * ep + right]; - } - } - - halo_q_direction8(common, count_left, count_right, quota_left, quota_right); - halo_q_direction8(common, count_right, count_left, quota_right, quota_left); - - if (common) - { - quota[b * ep + left] = quota_left; - quota[b * ep + right] = quota_right; - } -} - -__device__ __forceinline__ void halo_q_active_barrier(int threads) -{ - asm volatile("bar.sync 1, %0;" ::"r"(threads) : "memory"); -} - -__device__ __forceinline__ void scheduler_tail_barrier() -{ - asm volatile("bar.sync 3, 64;" ::: "memory"); -} - -/* Build order[] and host_mask[] with one lane per broadcast. - * - * This replaces serial insertion-sort and mask-construction loops while - * preserving the same stable order. - * - * order[]: lane b's position is the number of broadcasts that sort before it. - * The `j < lane` tiebreak makes this the same stable order the - * insertion sort produced (its shift test is a strict `>`), so the - * result is bit-identical, not merely equivalent. - * host_mask: the destination group is a contiguous rank range, so the mask is - * closed form; no loop over group members is needed. - * rank_group_mask handles the full EP32 word without a shift by32. - */ -__device__ __forceinline__ void halo_q_build_order_and_masks(int const* p_expert, int const* p_begin, int const* p_size, - int const* p_owner, int* order, int* host_mask, int broadcasts) -{ - int const lane = threadIdx.x & 31; - if (broadcasts > 32) - { - for (int b = lane; b < broadcasts; b += 32) - { - int const mine = p_expert[b]; - int position = 0; - for (int j = 0; j < broadcasts; ++j) - position += p_expert[j] < mine || (p_expert[j] == mine && j < b); - order[position] = b; - host_mask[b] = static_cast(rank_group_mask(p_begin[b], p_size[b]) | (1u << p_owner[b])); - } - return; - } - /* Read p_expert cooperatively and distribute values with shuffles. All - * lanes must participate before inactive lanes return; inactive lanes - * carry INT_MAX and never contribute to a real comparison. */ - bool const active = lane < broadcasts; - int const mine = active ? p_expert[lane] : INT_MAX; - int position = 0; - for (int j = 0; j < broadcasts; ++j) - { - int const other = __shfl_sync(0xffffffffu, mine, j); - position += (other < mine) || (other == mine && j < lane); - } - if (!active) - return; - order[position] = lane; - host_mask[lane] = static_cast(rank_group_mask(p_begin[lane], p_size[lane]) | (1u << p_owner[lane])); -} - -template -__device__ void run_halo_q(int const* counts, int* quota, int* plan, int broadcasts, int ep, int experts, int* ring, - int* order, int* host_mask) -{ - int const plan_ep = FixedEp > 0 ? FixedEp : ep; - int const tid = threadIdx.x; - int const warp = tid >> 5; - int const lane = tid & 31; - int* p_expert = plan_field(plan, kPlanExpert); - int* p_begin = plan_field(plan, kPlanGroupBegin); - int* p_size = plan_field(plan, kPlanGroupSize); - int* p_owner = plan_field(plan, kPlanOwner); - if (plan_ep <= 8 && broadcasts <= 8) - { - if (warp < plan_ep / 2) - { - if (warp == 0) - halo_q_build_order_and_masks(p_expert, p_begin, p_size, p_owner, order, host_mask, broadcasts); - halo_q_active_barrier((plan_ep / 2) * 32); - /* Loop-invariant broadcast identity, read once after the barrier - * that publishes order[] and host_mask[]. All four participating - * warps read the same lane-indexed entries. */ - int const my_b = lane < broadcasts ? order[lane] : -1; - int const my_hosts = my_b >= 0 ? host_mask[my_b] : 0; - int const my_expert = my_b >= 0 ? p_expert[my_b] : -1; - if (broadcasts >= 2) - { - int const ring_size = plan_ep - 1; - /* Use two unconditional repair rounds. Phases with no required repair - * are no-ops, so this preserves the conditional algorithm's result. */ - int const repair_rounds = 2; - for (int round = 0; round < repair_rounds; ++round) - { - for (int phase = 0; phase < plan_ep - 1; ++phase) - { - int right_index = warp == 0 ? 0 : plan_ep - 1 - warp; - int right_position = right_index - phase; - if (right_position < 0) - right_position += ring_size; - int right = 1 + right_position; - int left = 0; - if (warp != 0) - { - int left_position = warp - phase; - if (left_position < 0) - left_position += ring_size; - left = 1 + left_position; - } - halo_q_pair8(left, right, my_b, my_hosts, my_expert, counts, quota, plan_ep, experts); - if constexpr (FixedEp == 8) - { - /* Must stay unconditional. Omitting it after a - * round's last phase races twice: warp 0 enters the - * next round's phase 0 while warp 3 still writes the - * same quota[], and after the final phase warps 1-3 - * leave run_halo_q while warp 0 -- which only meets - * warp 8 at scheduler_tail_barrier -- walks into - * build_route_prefix and reads every quota[]. */ - halo_q_active_barrier(128); - } - else - { - halo_q_active_barrier((plan_ep / 2) * 32); - } - } - } - } - } - /* The unconditional repair rounds above replace the conditional scan. */ - return; - } - if (warp == 0) - halo_q_build_order_and_masks(p_expert, p_begin, p_size, p_owner, order, host_mask, broadcasts); - if (tid < plan_ep - 1) - ring[tid] = tid + 1; - __syncthreads(); - if (broadcasts < 2) - return; - - for (int round = 0; round < 2; ++round) - { - for (int phase = 0; phase < plan_ep - 1; ++phase) - { - if (warp < plan_ep / 2) - { - int left, right; - if (warp == 0) - { - left = 0; - right = ring[0]; - } - else - { - left = ring[warp]; - right = ring[plan_ep - 1 - warp]; - } - halo_q_repair(left, right, broadcasts, order, host_mask, counts, p_expert, quota, plan_ep, experts); - halo_q_repair(right, left, broadcasts, order, host_mask, counts, p_expert, quota, plan_ep, experts); - } - __syncthreads(); - if (tid == 0) - { - int last = ring[plan_ep - 2]; - for (int index = plan_ep - 2; index > 0; --index) - ring[index] = ring[index - 1]; - ring[0] = last; - } - __syncthreads(); - } - } - (void) lane; -} - -__device__ void color_and_publish_local(int* plan, int broadcasts, int ep, int helpers, int local_rank, int* used, - int* out_ids, int* out_levels, int* out_owners, int* status, Epoch epoch) -{ - int* p_expert = plan_field(plan, kPlanExpert); - int* p_level = plan_field(plan, kPlanLevel); - int* p_begin = plan_field(plan, kPlanGroupBegin); - int* p_size = plan_field(plan, kPlanGroupSize); - int* p_owner = plan_field(plan, kPlanOwner); - int* p_helper = plan_field(plan, kPlanHelper); - int const active_helpers = imin(helpers, kMaxBroadcasts); - int const words = (active_helpers + 31) / 32; - for (int rank = 0; rank < ep; ++rank) - for (int word = 0; word < words; ++word) - used[rank * kHelperMaskWords + word] = 0; - for (int level = 0; level < hierarchy_depth(ep); ++level) - { - for (int b = 0; b < broadcasts; ++b) - { - if (p_level[b] != level) - continue; - unsigned unavailable[kHelperMaskWords] = {0}; - for (int rank = p_begin[b]; rank < p_begin[b] + p_size[b]; ++rank) - for (int word = 0; word < words; ++word) - unavailable[word] |= static_cast(used[rank * kHelperMaskWords + word]); - int helper = -1; - for (int slot = 0; slot < active_helpers; ++slot) - if (!(unavailable[slot / 32] & (1u << (slot % 32)))) - { - helper = slot; - break; - } - if (helper < 0) - { - record_status(status + 4, epoch); - helper = 0; - } - p_helper[b] = helper; - for (int rank = p_begin[b]; rank < p_begin[b] + p_size[b]; ++rank) - used[rank * kHelperMaskWords + helper / 32] |= static_cast(1u << (helper % 32)); - } - } - for (int slot = 0; slot < helpers; ++slot) - { - out_ids[slot] = -1; - out_levels[slot] = -1; - out_owners[slot] = -1; - } - for (int b = 0; b < broadcasts; ++b) - { - if (!group_contains(local_rank, p_begin[b], p_size[b])) - continue; - int helper = p_helper[b]; - if (out_ids[helper] != -1) - record_status(status + 4, epoch); - out_ids[helper] = p_expert[b]; - out_levels[helper] = p_level[b]; - out_owners[helper] = p_owner[b]; - } -} - -__device__ void build_route_prefix(int const* counts, int const* quota, int* prefix, int const* plan, int broadcasts, - int ep, int experts, int local_rank, int participating_threads) -{ - int* p_expert = plan_field(const_cast(plan), kPlanExpert); - for (int b = threadIdx.x; b < broadcasts; b += participating_threads) - { - int expert = p_expert[b]; - int source_before = 0; - for (int rank = 0; rank < local_rank; ++rank) - { - int diagonal = imin(counts[rank * experts + expert], quota[b * ep + rank]); - source_before += counts[rank * experts + expert] - diagonal; - } - int source_count = counts[local_rank * experts + expert]; - int source_diagonal = imin(source_count, quota[b * ep + local_rank]); - int residual_supply = source_count - source_diagonal; - /* The loop used to run to destination == ep. route_destination only - * ever loads prefix[destination + step] under `candidate < ep`, so the - * ep entry is never read, and the accumulation at ep-1 fed nothing but - * that dead entry. Both are gone; entry 0 is still written because it - * costs one store and leaving it stale would depend on nobody ever - * reading it, which is a weaker guarantee than writing the right value. */ - int destination_before = 0; - for (int destination = 0; destination < ep; ++destination) - { - int overlap = destination_before - source_before; - overlap = imax(0, imin(residual_supply, overlap)); - int diagonal_prefix = local_rank < destination ? source_diagonal : 0; - prefix[b * (ep + 1) + destination] = diagonal_prefix + overlap; - if (destination + 1 < ep) - { - int diagonal = imin(counts[destination * experts + expert], quota[b * ep + destination]); - destination_before += quota[b * ep + destination] - diagonal; - } - } - } -} - -__device__ __forceinline__ int* worker_expert_lut(int experts) -{ - extern __shared__ int shared[]; - return shared + (blockDim.x >> 5) * (experts + 1); -} - -__device__ void prepare_worker_default_slots(int* expert_lut, int ep, int experts, int helpers) -{ - int const home = experts / ep; - int const local_slots = home + helpers; - for (int expert = threadIdx.x; expert < experts; expert += blockDim.x) - { - int owner = expert / home; - int local = expert - owner * home; - /* Negative values distinguish the precomputed cold slot from a - * nonnegative broadcast index. Slot zero therefore maps to -1. */ - expert_lut[expert] = -(owner * local_slots + local) - 1; - } - __syncthreads(); -} - -__device__ void prepare_worker_selected_experts(int* expert_lut, int const* plan) -{ - int const broadcasts = plan[1]; - int const* p_expert = plan + 2 + kPlanExpert * kMaxBroadcasts; - for (int b = threadIdx.x; b < broadcasts; b += blockDim.x) - expert_lut[p_expert[b]] = b; - __syncthreads(); -} - -__device__ bool wait_for_plan_epoch(int* marker, Epoch epoch, unsigned long long spin_cycles, int* status, int* ready) -{ - if (threadIdx.x == 0) - { - *ready = 1; - unsigned long long start = clock64(); - Epoch observed = load_relaxed_device(marker); - while (epoch_before(observed, epoch)) - { - if (clock64() - start > spin_cycles) - { - record_status(status + 2, epoch); - *ready = 0; - break; - } - observed = load_relaxed_device(marker); - } - if (!epoch_before(observed, epoch)) - (void) load_acquire_device(marker); - } - __syncthreads(); - return *ready != 0; -} - -/* Prefix entries are monotonic. Binary lifting returns the same number of - * thresholds <= ordinal as the old linear scan, including repeated entries. */ -__device__ __forceinline__ int route_destination(int const* prefix, int ordinal, int ep) -{ - int destination = 0; -#pragma unroll - for (int step = 16; step; step >>= 1) - { - int candidate = destination + step; - if (candidate < ep && ordinal >= prefix[candidate]) - destination = candidate; - } - return destination; -} - -/* Once HALO-M publishes its selected experts, workers can finish every cold - * route while CTA0 runs HALO-Q. Selected route indices are compacted into - * disjoint slices of route_aux; their ordinal values remain in output. */ -__device__ void pre_materialize_and_compact_worker_routes(int const* routes, int* output, int* route_aux, int* status, - int route_count, int experts, int worker, int workers, Epoch epoch) -{ - extern __shared__ int shared[]; - int const tid = threadIdx.x; - int const lane = tid & 31; - int const warp = tid >> 5; - int const warps = blockDim.x >> 5; - int const bins = experts + 1; - int* expert_lut = worker_expert_lut(experts); - int* hot_indices = route_aux + 2 + workers * bins; - - int const worker_span = (route_count + workers - 1) / workers; - int const worker_begin = worker * worker_span; - int const worker_end = imin(route_count, worker_begin + worker_span); - int const warp_span = ((worker_span + warps * 32 - 1) / (warps * 32)) * 32; - int const begin = worker_begin + warp * warp_span; - int const end = imin(worker_end, begin + warp_span); - int hot_count = 0; - for (int base = begin; base < end; base += 32) - { - int index = base + lane; - bool active = index < end; - int expert = active ? routes[index] : -1; - int slot = -1; - bool hot = false; - if (active && expert >= 0 && expert < experts) - { - int entry = expert_lut[expert]; - hot = entry >= 0; - if (!hot) - slot = -entry - 1; - } - else if (active && expert != -1) - { - record_status(status + 3, epoch); - } - unsigned hot_mask = __ballot_sync(0xffffffffu, hot); - if (hot) - { - int offset = __popc(hot_mask & ((1u << lane) - 1u)); - hot_indices[begin + hot_count + offset] = index; - } - else if (active) - { - output[index] = slot; - } - hot_count += __popc(hot_mask); - } - if (lane == 0) - shared[warp * bins + experts] = hot_count; -} - -/* The final epoch only has to resolve selected routes. Stable ordinal - * cursors remain in the per-warp shared histogram from the initial pass. */ -__device__ void materialize_worker_hot_routes(int const* routes, int* output, int const* route_prefix, int const* plan, - int* route_aux, int* status, int route_count, int ep, int experts, int helpers, int worker, int workers, - Epoch epoch) -{ - extern __shared__ int shared[]; - int const lane = threadIdx.x & 31; - int const warp = threadIdx.x >> 5; - int const warps = blockDim.x >> 5; - int const bins = experts + 1; - int* expert_lut = worker_expert_lut(experts); - int* hot_indices = route_aux + 2 + workers * bins; - int const home = experts / ep; - int const local_slots = home + helpers; - int const* p_begin = plan + 2 + kPlanGroupBegin * kMaxBroadcasts; - int const* p_size = plan + 2 + kPlanGroupSize * kMaxBroadcasts; - int const* p_owner = plan + 2 + kPlanOwner * kMaxBroadcasts; - int const* p_helper = plan + 2 + kPlanHelper * kMaxBroadcasts; - - int const worker_span = (route_count + workers - 1) / workers; - int const worker_begin = worker * worker_span; - int const worker_end = imin(route_count, worker_begin + worker_span); - int const warp_span = ((worker_span + warps * 32 - 1) / (warps * 32)) * 32; - int const begin = worker_begin + warp * warp_span; - int const end = imin(worker_end, begin + warp_span); - int hot_count = shared[warp * bins + experts]; - hot_count = imin(hot_count, imax(0, end - begin)); - for (int offset = lane; offset < hot_count; offset += 32) - { - int index = hot_indices[begin + offset]; - int expert = routes[index]; - int b = expert_lut[expert]; - int ordinal = output[index] + shared[warp * bins + expert]; - int const* prefix = route_prefix + b * (ep + 1); - int destination = route_destination(prefix, ordinal, ep); - int slot = -1; - if (group_contains(destination, p_begin[b], p_size[b])) - { - slot = destination * local_slots + home + p_helper[b]; - } - else if (destination == p_owner[b]) - { - slot = destination * local_slots + expert - p_owner[b] * home; - } - else - { - record_status(status + 5, epoch); - } - output[index] = slot; - } -} - -__device__ void materialize_worker_routes(int const* routes, int* output, int const* route_prefix, int const* plan, - int* status, int route_count, int ep, int experts, int helpers, int worker, int workers, Epoch epoch) -{ - extern __shared__ int shared[]; - int const tid = threadIdx.x; - int const lane = tid & 31; - int const warp = tid >> 5; - int const warps = blockDim.x >> 5; - int const bins = experts + 1; - int* expert_lut = worker_expert_lut(experts); - int const home = experts / ep; - int const local_slots = home + helpers; - int const* p_begin = plan + 2 + kPlanGroupBegin * kMaxBroadcasts; - int const* p_size = plan + 2 + kPlanGroupSize * kMaxBroadcasts; - int const* p_owner = plan + 2 + kPlanOwner * kMaxBroadcasts; - int const* p_helper = plan + 2 + kPlanHelper * kMaxBroadcasts; - - int const worker_span = (route_count + workers - 1) / workers; - int const worker_begin = worker * worker_span; - int const worker_end = imin(route_count, worker_begin + worker_span); - int const warp_span = ((worker_span + warps * 32 - 1) / (warps * 32)) * 32; - int const begin = worker_begin + warp * warp_span; - int const end = imin(worker_end, begin + warp_span); - for (int index = begin + lane; index < end; index += 32) - { - int expert = routes[index]; - int slot = -1; - if (expert >= 0 && expert < experts) - { - int entry = expert_lut[expert]; - if (entry < 0) - { - slot = -entry - 1; - } - else - { - int b = entry; - int ordinal = output[index] + shared[warp * bins + expert]; - int const* prefix = route_prefix + b * (ep + 1); - int destination = route_destination(prefix, ordinal, ep); - if (group_contains(destination, p_begin[b], p_size[b])) - { - slot = destination * local_slots + home + p_helper[b]; - } - else if (destination == p_owner[b]) - { - slot = destination * local_slots + expert - p_owner[b] * home; - } - else - { - record_status(status + 5, epoch); - } - } - } - else if (expert != -1) - { - record_status(status + 3, epoch); - } - output[index] = slot; - } -} template __global__ __launch_bounds__(512, 1) void halo_q_scheduler_kernel(int const* routes, int* out_slots, int* out_ids, @@ -1784,192 +42,10 @@ __global__ __launch_bounds__(512, 1) void halo_q_scheduler_kernel(int const* rou int runtime_helpers, int local_rank, int route_capacity, int ctas, int enable_pdl, unsigned long long spin_cycles, int plan_abi_version, int plan_channel_words, int route_features, int capacity, int route_count) { - int const ep = FixedEp > 0 ? FixedEp : runtime_ep; - int const experts = FixedExperts > 0 ? FixedExperts : runtime_experts; - int const helpers = FixedHelpers > 0 ? FixedHelpers : runtime_helpers; - (void) plan_channel_words; // Validated by the host launch wrapper. - /* The scheduler grid is deliberately residency-safe (one CTA per SM and - * CTA count <= SM count), so signal as soon as every CTA starts. The - * independent same-stream quantizer may then occupy all unused SMs while - * these CTAs retain their own SMs for the complete scheduler lifetime. */ - if (enable_pdl && threadIdx.x == 0) - cudaTriggerProgrammaticLaunchCompletion(); - // Only the caller's valid prefix is readable. Keep the public output at - // its configured capacity, clearing stale routes inside this same kernel. - // Empty ranks still execute every collective; no route load occurs at N=0. - for (int index = route_count + blockIdx.x * blockDim.x + threadIdx.x; index < route_capacity; - index += gridDim.x * blockDim.x) - out_slots[index] = -1; - local_histogram_pass(routes, partial, route_count, experts, ctas); - Epoch epoch; - if (!grid_rendezvous(grid_sync, ctas, spin_cycles, status, &epoch)) - return; - - if (blockIdx.x == 0) - { - extern __shared__ int shared[]; - int offset = ep * experts; - int* counts = shared; - int* histogram = shared + offset; - offset += experts; - int* load = shared + offset; - offset += ep; - // Reuse the exact capacity already validated/computed by launch(). - int* quota = shared + offset; - offset += capacity * ep; - int* selected_flags = shared + offset; - offset += experts; - int* recv_used = shared + offset; - offset += ep * kHelperMaskWords; - int* send_used = shared + offset; - offset += ep; - int* next_index = shared + offset; - offset += ep; - int* frozen = shared + offset; - offset += ep; - int* ring = shared + offset; - offset += ep; - int* order = shared + offset; - offset += capacity; - int* host_mask = shared + offset; - - bool exchange_ok; - if (ep == 8) - exchange_ok = exchange_counts<8>( - peer_bases, partial, shared, ep, experts, local_rank, ctas, epoch, spin_cycles, status); - else - exchange_ok = exchange_counts<0>( - peer_bases, partial, shared, ep, experts, local_rank, ctas, epoch, spin_cycles, status); - if (!exchange_ok) - return; - /* EP8 keeps only its candidate masses; other EPs retain the dense - * histogram required by the generic planner path. */ - if (ep == 8 && helpers == 4 && experts >= 32) - { - prepare_halo_m_sparse8(counts, histogram, load, selected_flags, experts); - } - else - { - prepare_halo_m<0, 0>(counts, histogram, load, selected_flags, recv_used, send_used, next_index, frozen, ep, - experts, helpers); - } - - int broadcasts = 0; - if (threadIdx.x < (ep <= 8 ? 8 : 32)) - { - // These short-circuit bounds keep helpers*ep <= 120. - if (ep <= 8 && helpers <= 15 && helpers * ep <= experts) - { - if (ep == 8 && helpers == 4) - { - broadcasts = plan_halo_m_register8<8, 4>(histogram, load, quota, selected_flags, plan, ep, helpers); - } - else if (ep == 4 && helpers == 4) - { - broadcasts = plan_halo_m_register8<4, 4>(histogram, load, quota, selected_flags, plan, ep, helpers); - } - else - { - broadcasts = plan_halo_m_register8<0, 0>(histogram, load, quota, selected_flags, plan, ep, helpers); - } - } - else - { - broadcasts = plan_halo_m_warp(histogram, load, quota, selected_flags, recv_used, send_used, next_index, - frozen, plan, ep, experts, helpers); - } - if (threadIdx.x == 0) - plan[1] = broadcasts; - } - __syncthreads(); - broadcasts = plan[1]; - /* p_expert and plan[1] are produced by thread zero. Publish this - * placement-only epoch before HALO-Q so workers can construct their - * selected-expert LUT while quota repair is still running. */ - if (ctas > 1 && threadIdx.x == 0) - { - store_release_device(plan + 0, epoch); - } - /* Small-EP repair excludes warp8 and can overlap coloring. Generic - * repair first meets the coloring warp at its full-block barrier. */ - if (threadIdx.x == 8 * 32) - color_and_publish_local( - plan, broadcasts, ep, helpers, local_rank, recv_used, out_ids, out_levels, out_owners, status, epoch); - if (ep == 8) - run_halo_q<8>(counts, quota, plan, broadcasts, ep, experts, ring, order, host_mask); - else - run_halo_q<0>(counts, quota, plan, broadcasts, ep, experts, ring, order, host_mask); - if (ctas > 1) - { - int const warp = threadIdx.x >> 5; - /* build_route_prefix is independent of warp 8's outputs. Defer their - * rendezvous until both paths reach their first shared consumer so - * prefix construction and coloring can overlap safely. */ - if (warp == 0) - { - build_route_prefix(counts, quota, route_prefix, plan, broadcasts, ep, experts, local_rank, 32); - __syncwarp(); - } - if (warp == 0 || warp == 8) - scheduler_tail_barrier(); - if (warp == 0) - { - /* Release the worker CTAs here rather than after the channel - * publish. They gate on grid_sync+2 and consume route_prefix, - * which is complete above; plan_channel is consumed only by the - * copy stack, which acquires on the channel's own epoch word - * written with store_release_system below. Publishing after - * the release lets the workers coloring pass -- the work that - * actually determines when the kernel ends -- overlap the whole - * publish instead of waiting behind it. */ - if (threadIdx.x == 0) - store_release_device(grid_sync + 2, epoch); - sami_copy::publish_plan_channel_warp(plan_channel, out_ids, out_levels, out_owners, - copy_placement_view(plan), helpers, ep, local_rank, experts, plan_abi_version, route_features); - __syncwarp(); - } - return; - } - if (threadIdx.x == 0) - store_release_device(plan + 0, epoch); - __syncthreads(); - build_route_prefix(counts, quota, route_prefix, plan, broadcasts, ep, experts, local_rank, blockDim.x); - __syncthreads(); - sami_copy::publish_plan_channel_warp(plan_channel, out_ids, out_levels, out_owners, copy_placement_view(plan), - helpers, ep, local_rank, experts, plan_abi_version, route_features); - /* The block barrier orders route-prefix/helper writes before release. */ - __syncthreads(); - } - else - { - int const worker = blockIdx.x - 1; - int const workers = ctas - 1; - if (!stable_worker_ordinal_pass( - routes, out_slots, route_aux, status, route_count, experts, worker, workers, spin_cycles)) - return; - int* expert_lut = worker_expert_lut(experts); - prepare_worker_default_slots(expert_lut, ep, experts, helpers); - __shared__ int plan_ready; - if (!wait_for_plan_epoch(plan + 0, epoch, spin_cycles, status, &plan_ready)) - return; - prepare_worker_selected_experts(expert_lut, plan); - pre_materialize_and_compact_worker_routes( - routes, out_slots, route_aux, status, route_count, experts, worker, workers, epoch); - if (!wait_for_plan_epoch(grid_sync + 2, epoch, spin_cycles, status, &plan_ready)) - return; - materialize_worker_hot_routes(routes, out_slots, route_prefix, plan, route_aux, status, route_count, ep, - experts, helpers, worker, workers, epoch); - return; - } - - /* CTA-count one remains a correct serial fallback. */ - if (!stable_worker_ordinal_pass(routes, out_slots, route_aux, status, route_count, experts, 0, 1, spin_cycles)) - return; - int* expert_lut = worker_expert_lut(experts); - prepare_worker_default_slots(expert_lut, ep, experts, helpers); - prepare_worker_selected_experts(expert_lut, plan); - materialize_worker_routes( - routes, out_slots, route_prefix, plan, status, route_count, ep, experts, helpers, 0, 1, epoch); + halo_q_scheduler_body(routes, out_slots, out_ids, out_levels, out_owners, + peer_bases, status, partial, route_aux, grid_sync, plan, route_prefix, plan_channel, runtime_ep, + runtime_experts, runtime_helpers, local_rank, route_capacity, ctas, enable_pdl, spin_cycles, plan_abi_version, + plan_channel_words, route_features, capacity, route_count); } bool validPlanChannelLaunch(std::uintptr_t pointer, int abiVersion, int words, int features, int ep, int helpers) @@ -2010,7 +86,7 @@ cudaError_t ensureDynamicSharedMemory(void const* kernel, int sharedBytes) } // namespace -void invokeMoeRebalanceHaloQ(MoeRebalanceHaloQParams const& params, cudaStream_t stream) +void checkMoeRebalanceHaloQParams(MoeRebalanceHaloQParams const& params) { TLLM_CHECK_WITH_INFO((params.routes != nullptr || params.validRouteCount == 0) && params.outSlots != nullptr && params.outIds != nullptr && params.outLevels != nullptr && params.outOwners != nullptr @@ -2028,15 +104,16 @@ void invokeMoeRebalanceHaloQ(MoeRebalanceHaloQParams const& params, cudaStream_t validPlanChannelLaunch(reinterpret_cast(params.planChannel), params.planAbiVersion, params.planChannelWords, params.routeFeatures, params.ep, params.helpers), "Invalid HALO-Q PlanChannel launch ABI"); +} + +void invokeMoeRebalanceHaloQ(MoeRebalanceHaloQParams const& params, cudaStream_t stream) +{ + checkMoeRebalanceHaloQParams(params); - int const warps = params.threads / 32; - int const passInts = warps * (params.experts + 1); - int const workerInts = passInts + params.experts; int capacity = broadcast_capacity(params.ep, params.experts, params.helpers); - int const plannerInts = params.ep * params.experts + params.experts + params.ep + capacity * params.ep - + params.experts + (kHelperMaskWords + 4) * params.ep + 2 * capacity; std::size_t const sharedBytes - = static_cast(workerInts > plannerInts ? workerInts : plannerInts) * sizeof(int); + = static_cast(halo_q_shared_ints(params.ep, params.experts, capacity, params.threads)) + * sizeof(int); void const* kernel = reinterpret_cast(&halo_q_scheduler_kernel<0, 0, 0>); if (params.experts == 384 && params.helpers == 4) diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h index d0215ab74617..5d0d10a62bb6 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h @@ -58,6 +58,9 @@ struct MoeRebalanceHaloQParams int validRouteCount; }; +// Throws unless params describe a valid scheduler launch. +void checkMoeRebalanceHaloQParams(MoeRebalanceHaloQParams const& params); + void invokeMoeRebalanceHaloQ(MoeRebalanceHaloQParams const& params, cudaStream_t stream); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh new file mode 100644 index 000000000000..9f4733972f65 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh @@ -0,0 +1,1994 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/* HALO-Q scheduler device code. The standalone scheduler kernel and the fused + * scheduler + in-switch copy kernel both run halo_q_scheduler_body. */ +#ifndef TRTLLM_KERNELS_MOE_REBALANCE_HALO_Q_KERNEL_CUH +#define TRTLLM_KERNELS_MOE_REBALANCE_HALO_Q_KERNEL_CUH + +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlanPolicy.h" + +#include +#include +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ +constexpr int kMaxEp = 32; +constexpr int kMaxExperts = 384; +constexpr int kMaxBroadcasts = kMaxExperts; +constexpr int kHelperMaskWords = (kMaxBroadcasts + 31) / 32; +constexpr int kMinGroup = 4; +constexpr int kSymPayloadOffset = 64; +using Epoch = std::uint32_t; + +enum PlanField +{ + kPlanExpert = 0, + kPlanLevel = 1, + kPlanGroupBegin = 2, + kPlanGroupSize = 3, + kPlanOwner = 4, + kPlanHelper = 5, +}; + +// Adapt placement storage to the copy layer without exposing scheduler +// internals. +__device__ __forceinline__ sami_copy::PlacementView copy_placement_view(int const* plan) +{ + return {plan[1], kMaxBroadcasts, plan + 2 + kPlanExpert * kMaxBroadcasts, plan + 2 + kPlanLevel * kMaxBroadcasts, + plan + 2 + kPlanGroupBegin * kMaxBroadcasts, plan + 2 + kPlanGroupSize * kMaxBroadcasts, + plan + 2 + kPlanOwner * kMaxBroadcasts, plan + 2 + kPlanHelper * kMaxBroadcasts}; +} + +__device__ __forceinline__ int imin(int a, int b) +{ + return a < b ? a : b; +} + +__device__ __forceinline__ int imax(int a, int b) +{ + return a > b ? a : b; +} + +__device__ __forceinline__ Epoch fetch_add_acq_rel_device(int* address, Epoch value) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + return atom.fetch_add(value, cuda::memory_order_acq_rel); +} + +__device__ __forceinline__ Epoch load_acquire_device(int* address) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + return atom.load(cuda::memory_order_acquire); +} + +__device__ __forceinline__ Epoch load_relaxed_device(int* address) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + return atom.load(cuda::memory_order_relaxed); +} + +__device__ __forceinline__ void store_relaxed_device(int* address, Epoch value) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + atom.store(value, cuda::memory_order_relaxed); +} + +__device__ __forceinline__ void store_release_device(int* address, Epoch value) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + atom.store(value, cuda::memory_order_release); +} + +__device__ __forceinline__ Epoch load_acquire_system(int* address) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + return atom.load(cuda::memory_order_acquire); +} + +__device__ __forceinline__ Epoch load_relaxed_system(int* address) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + return atom.load(cuda::memory_order_relaxed); +} + +__device__ __forceinline__ void store_release_system(int* address, Epoch value) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + atom.store(value, cuda::memory_order_release); +} + +__device__ __forceinline__ bool epoch_before(Epoch observed, Epoch expected) +{ + // RFC-1982-style serial comparison: all arithmetic remains unsigned and + // is unambiguous while peers differ by fewer than half the epoch space. + return observed != expected && expected - observed < (Epoch{1} << 31); +} + +__device__ __forceinline__ void record_status(int* address, Epoch epoch) +{ + // Zero is reserved for "never failed". Preserve a nonzero sticky error at + // the one modular epoch whose bit pattern is zero. + Epoch const value = epoch == 0 ? ~Epoch{0} : epoch; + atomicExch(reinterpret_cast(address), value); +} + +__device__ __forceinline__ int* plan_field(int* plan, int field) +{ + return plan + 2 + field * kMaxBroadcasts; +} + +__device__ __forceinline__ bool group_contains(int rank, int begin, int size) +{ + return rank >= begin && rank < begin + size; +} + +__host__ __device__ __forceinline__ unsigned rank_group_mask(int begin, int size) +{ + return size == 32 ? 0xffffffffu : ((1u << size) - 1u) << begin; +} + +__host__ __device__ __forceinline__ int hierarchy_depth(int ep) +{ + int depth = 1; + for (int level = 1; (1 << level) <= ep; ++level) + { + int groups = 1 << level; + if (ep % groups || ep / groups < kMinGroup) + break; + ++depth; + } + return depth; +} + +/* Distinct experts and total receive capacity independently bound broadcasts. + * Actual S is retained in the public outputs; only scratch uses this bound. */ +__host__ __device__ __forceinline__ int broadcast_capacity(int ep, int experts, int helpers) +{ + int const bounded_helpers = helpers < experts ? helpers : experts; + /* Every hierarchy level divides ep exactly. This cancels two integer + * divisions without changing capacity, including non-power-of-two EP. */ + int const groups = 1 << (hierarchy_depth(ep) - 1); + int const receive_bound = groups * bounded_helpers; + return receive_bound < experts ? receive_bound : experts; +} + +__device__ bool grid_rendezvous( + int* sync, int blocks, unsigned long long spin_cycles, int* status, Epoch* completedEpoch) +{ + __shared__ Epoch generation; + __shared__ int success; + /* The block barrier orders every writer's stores before thread zero's + * device-scope release advertises this CTA. */ + __syncthreads(); + if (threadIdx.x == 0) + { + success = 1; + Epoch const current = load_relaxed_device(sync + 1); + Epoch const target = current + 1; + generation = target; + Epoch const old = fetch_add_acq_rel_device(sync, 1); + if (old == static_cast(blocks - 1)) + { + // The last CTA resets arrivals before publishing the next modular + // epoch. Stream ordering prevents a new launch from entering until + // every CTA in this launch has observed the release. + store_relaxed_device(sync, 0); + store_release_device(sync + 1, target); + } + else if (old >= static_cast(blocks)) + { + record_status(status + 1, target); + success = 0; + } + else + { + unsigned long long const start = clock64(); + Epoch observed = load_relaxed_device(sync + 1); + while (epoch_before(observed, target)) + { + if (clock64() - start > spin_cycles) + { + record_status(status + 1, target); + success = 0; + break; + } + observed = load_relaxed_device(sync + 1); + } + if (!epoch_before(observed, target)) + (void) load_acquire_device(sync + 1); + } + } + __syncthreads(); + *completedEpoch = generation; + return success != 0; +} + +__device__ void local_histogram_pass(int const* routes, int* partial, int route_count, int expert_count, int ctas) +{ + extern __shared__ int shared[]; + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const warps = blockDim.x >> 5; + int const bins = expert_count + 1; + int const shards = warps > 8 ? 8 : warps; + int const shared_count = shards * bins; + for (int index = tid; index < shared_count; index += blockDim.x) + shared[index] = 0; + __syncthreads(); + + int const cta_span = (route_count + ctas - 1) / ctas; + int const cta_begin = blockIdx.x * cta_span; + int const cta_end = imin(route_count, cta_begin + cta_span); + int const warp_span = ((cta_span + warps * 32 - 1) / (warps * 32)) * 32; + int const begin = cta_begin + warp * warp_span; + int const end = imin(cta_end, begin + warp_span); + /* Issue a small load group before consuming it to expose memory-level + * parallelism. The binning is unchanged, and atomicAdd regrouping is + * safe because integer addition commutes. */ + constexpr int kBatch = 4; + for (int base_index = begin; base_index < end; base_index += 32 * kBatch) + { + int expert[kBatch]; +#pragma unroll + for (int step = 0; step < kBatch; ++step) + { + int const index = base_index + step * 32 + lane; + expert[step] = index < end ? routes[index] : -1; + } +#pragma unroll + for (int step = 0; step < kBatch; ++step) + { + int const index = base_index + step * 32 + lane; + if (index < end) + { + int const value = expert[step]; + int const bin = (value >= 0 && value < expert_count) ? value : expert_count; + atomicAdd(shared + (warp & 7) * bins + bin, 1); + } + } + } + __syncthreads(); + + for (int expert = tid; expert < expert_count; expert += blockDim.x) + { + int count = 0; + for (int shard = 0; shard < shards; ++shard) + count += shared[shard * bins + expert]; + partial[blockIdx.x * expert_count + expert] = count; + } +} + +/* Re-rank all local routes on the C-1 worker CTAs while CTA0 performs the + * cross-rank exchange and planner. route_aux[0:2] is a reusable worker + * rendezvous; the following NWORK*(E+1) entries hold per-worker totals. + * Per-warp cursors remain in this CTA's shared memory until materialization. */ +__device__ bool stable_worker_ordinal_pass(int const* routes, int* ordinal, int* route_aux, int* status, + int route_count, int expert_count, int worker, int workers, unsigned long long spin_cycles) +{ + extern __shared__ int shared[]; + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const warps = blockDim.x >> 5; + int const bins = expert_count + 1; + int const shared_count = warps * bins; + for (int index = tid; index < shared_count; index += blockDim.x) + shared[index] = 0; + __syncthreads(); + + int const worker_span = (route_count + workers - 1) / workers; + int const worker_begin = worker * worker_span; + int const worker_end = imin(route_count, worker_begin + worker_span); + int const warp_span = ((worker_span + warps * 32 - 1) / (warps * 32)) * 32; + int const begin = worker_begin + warp * warp_span; + int const end = imin(worker_end, begin + warp_span); + for (int base_index = begin; base_index < end; base_index += 32) + { + int const index = base_index + lane; + bool const active = index < end; + int const expert = active ? routes[index] : -1; + // Padding has its own key so it cannot contribute to real/invalid bins. + int const bin = !active ? bins : ((expert >= 0 && expert < expert_count) ? expert : expert_count); + unsigned const peers = __match_any_sync(0xffffffffu, bin); + int const leader = __ffs(peers) - 1; + int const rank = __popc(peers & ((1u << lane) - 1u)); + int base = 0; + if (active && lane == leader) + base = atomicAdd(shared + warp * bins + bin, __popc(peers)); + // One collective for the entire warp; each lane selects its own leader. + base = __shfl_sync(0xffffffffu, base, leader); + if (active) + ordinal[index] = base + rank; + } + __syncthreads(); + + int* worker_totals = route_aux + 2; + for (int bin = tid; bin < bins; bin += blockDim.x) + { + int total = 0; + for (int local_warp = 0; local_warp < warps; ++local_warp) + total += shared[local_warp * bins + bin]; + worker_totals[worker * bins + bin] = total; + } + + /* The first all-CTA rendezvous guarantees every worker is resident, so a + * worker-only global barrier cannot wait on an unscheduled CTA. */ + Epoch workerEpoch; + if (!grid_rendezvous(route_aux, workers, spin_cycles, status, &workerEpoch)) + return false; + + for (int bin = tid; bin < bins; bin += blockDim.x) + { + int cursor = 0; + for (int previous = 0; previous < worker; ++previous) + cursor += worker_totals[previous * bins + bin]; + for (int local_warp = 0; local_warp < warps; ++local_warp) + { + int index = local_warp * bins + bin; + int count = shared[index]; + shared[index] = cursor; + cursor += count; + } + } + __syncthreads(); + return true; +} + +template +__device__ bool exchange_counts(unsigned long long const* peer_bases, int const* partial, int* shared, int ep, + int expert_count, int local_rank, int ctas, Epoch epoch, unsigned long long spin_cycles, int* status) +{ + int const exchange_ep = FixedEp > 0 ? FixedEp : ep; + int const tid = threadIdx.x; + __shared__ int wait_failed; + if (tid == 0) + wait_failed = 0; + /* Keep the simple accumulation loop; the compiler provides the required + * pipelining without explicit batching. */ + for (int expert = tid; expert < expert_count; expert += blockDim.x) + { + int count = 0; + for (int cta = 0; cta < ctas; ++cta) + count += partial[cta * expert_count + expert]; + shared[expert] = count; + } + __syncthreads(); + + int const parity = epoch & 1; + /* One warp owns each destination row. Its warp barrier orders every + * lane's stores before lane zero's system-scope release publishes the + * complete row to the peer acquire. */ + int const warp = tid >> 5; + int const lane = tid & 31; + for (int destination = warp; destination < exchange_ep; destination += blockDim.x >> 5) + { + int* peer = reinterpret_cast(peer_bases[destination]); + int* row = peer + kSymPayloadOffset + (parity * exchange_ep + local_rank) * expert_count; + /* Keep scalar remote stores for the publication path; the local mirror + * copy below uses vector transfers when alignment permits. */ + for (int expert = lane; expert < expert_count; expert += 32) + row[expert] = shared[expert]; + /* Warp synchronization orders every lane's row stores before lane + * zero resumes; its system-scope release then publishes that complete + * transitive happens-before set to the peer acquire. */ + __syncwarp(); + if (lane == 0) + store_release_system(peer + local_rank, epoch); + } + __syncthreads(); + + int* local = reinterpret_cast(peer_bases[local_rank]); + if (tid < exchange_ep) + { + unsigned long long start = clock64(); + Epoch observed = load_relaxed_system(local + tid); + while (epoch_before(observed, epoch)) + { + if (clock64() - start > spin_cycles) + { + record_status(status + 0, epoch); + atomicExch(&wait_failed, 1); + break; + } + observed = load_relaxed_system(local + tid); + } + if (!epoch_before(observed, epoch)) + (void) load_acquire_system(local + tid); + } + __syncthreads(); + if (wait_failed) + return false; + int const* payload = local + kSymPayloadOffset + parity * exchange_ep * expert_count; + int const total = exchange_ep * expert_count; + /* Vectorize the local mirror copy when shape and alignment permit. */ + if ((total & 3) == 0 && (reinterpret_cast(payload) & 15) == 0) + { + int4 const* source = reinterpret_cast(payload); + int4* destination = reinterpret_cast(shared); + int const vectors = total >> 2; + for (int vector = tid; vector < vectors; vector += blockDim.x) + destination[vector] = source[vector]; + } + else + { + for (int index = tid; index < total; index += blockDim.x) + shared[index] = payload[index]; + } + __syncthreads(); + return true; +} + +__device__ void prepare_halo_m_sparse8(int const* counts, int* histogram, int* load, int* candidates, int experts) +{ + int const tid = threadIdx.x; + int const warp = tid >> 5; + int const lane = tid & 31; + if (warp >= 8) + return; + int const home = experts / 8; + int local_index[2]; + int local_mass[2]; + int rank_load = 0; +#pragma unroll + for (int round = 0; round < 2; ++round) + { + int local = lane + 32 * round; + int mass = -1; + if (local < home) + { + mass = 0; + int expert = warp * home + local; +#pragma unroll + for (int source = 0; source < 8; ++source) + mass += counts[source * experts + expert]; + rank_load += mass; + } + local_index[round] = local; + local_mass[round] = mass; + } + rank_load = __reduce_add_sync(0xffffffffu, rank_load); + if (lane == 0) + load[warp] = rank_load; + for (int slot = 0; slot < 4; ++slot) + { + int best_mass = -1; + int best_local = INT_MAX; +#pragma unroll + for (int round = 0; round < 2; ++round) + { + int mass = local_mass[round]; + int local = local_index[round]; + if (mass > best_mass || (mass == best_mass && local < best_local)) + { + best_mass = mass; + best_local = local; + } + } + int warp_mass = __reduce_max_sync(0xffffffffu, best_mass); + int winner = __reduce_min_sync(0xffffffffu, best_mass == warp_mass ? best_local : INT_MAX); + if (lane == 0) + { + int expert = warp * home + winner; + candidates[warp * 4 + slot] = expert; + histogram[expert] = warp_mass; + } +#pragma unroll + for (int round = 0; round < 2; ++round) + if (local_index[round] == winner) + local_mass[round] = -1; + } + asm volatile("bar.sync 4, 256;" ::: "memory"); +} + +/* Build HALO-M's dense histogram, rank loads and per-owner hot lists with one + * warp per rank. The former thread-0 implementation dominated the kernel. */ +template +__device__ void prepare_halo_m(int const* counts, int* histogram, int* load, int* candidates, int* recv_used, + int* send_used, int* next_index, int* frozen, int ep, int experts, int helpers) +{ + int const prepare_ep = FixedEp > 0 ? FixedEp : ep; + int const prepare_helpers = imin(FixedHelpers > 0 ? FixedHelpers : helpers, experts / prepare_ep); + int const tid = threadIdx.x; + int const warp = tid >> 5; + int const lane = tid & 31; + int const home = experts / prepare_ep; + for (int expert = tid; expert < experts; expert += blockDim.x) + { + int mass = 0; +#pragma unroll + for (int rank = 0; rank < prepare_ep; ++rank) + mass += counts[rank * experts + expert]; + histogram[expert] = mass; + } + __syncthreads(); + + for (int owner = warp; owner < prepare_ep; owner += blockDim.x >> 5) + { + int rank_load = 0; + for (int local = lane; local < home; local += 32) + rank_load += histogram[owner * home + local]; + rank_load = __reduce_add_sync(0xffffffffu, rank_load); + if (lane == 0) + { + load[owner] = rank_load; + recv_used[owner] = 0; + send_used[owner] = 0; + next_index[owner] = 0; + frozen[owner] = 0; + } + + /* EP2 is the largest home: 192 experts, six entries per lane. */ + constexpr int rounds = FixedEp == 8 ? 2 : 6; + int local_index[rounds]; + int local_mass[rounds]; +#pragma unroll + for (int round = 0; round < rounds; ++round) + { + int local = lane + 32 * round; + local_index[round] = local; + local_mass[round] = local < home ? histogram[owner * home + local] : -1; + } + for (int slot = 0; slot < prepare_helpers; ++slot) + { + int best_mass = -1; + int best_local = INT_MAX; +#pragma unroll + for (int round = 0; round < rounds; ++round) + { + int mass = local_mass[round]; + int local = local_index[round]; + if (mass > best_mass || (mass == best_mass && local < best_local)) + { + best_mass = mass; + best_local = local; + } + } + int warp_mass = __reduce_max_sync(0xffffffffu, best_mass); + int winner = __reduce_min_sync(0xffffffffu, best_mass == warp_mass ? best_local : INT_MAX); + if (lane == 0) + candidates[owner * prepare_helpers + slot] = owner * home + winner; +#pragma unroll + for (int round = 0; round < rounds; ++round) + if (local_index[round] == winner) + local_mass[round] = -1; + } + } + __syncthreads(); +} + +/* EP<=8 fast path: every participating lane owns one rank's complete + * mutable HALO-M state. Group and destination decisions use only shuffles, + * avoiding repeated shared-memory rank-state traffic and warp barriers. */ +template +__device__ __forceinline__ int plan_halo_m_register8_impl( + int const* histogram, int lane_load, int total, int* quota, int const* candidates, int* plan, int ep, int helpers) +{ + int const plan_ep = FixedEp > 0 ? FixedEp : ep; + int const plan_helpers = FixedHelpers > 0 ? FixedHelpers : helpers; + int const lane = threadIdx.x & 31; + constexpr unsigned mask = 0xffu; + int lane_recv = 0; + int lane_send = 0; + int lane_next = 0; + int lane_frozen = 0; + + int const target = total / plan_ep + (total % plan_ep != 0); + int const depth = hierarchy_depth(plan_ep); + int* p_expert = plan_field(plan, kPlanExpert); + int* p_level = plan_field(plan, kPlanLevel); + int* p_begin = plan_field(plan, kPlanGroupBegin); + int* p_size = plan_field(plan, kPlanGroupSize); + int* p_owner = plan_field(plan, kPlanOwner); + int* p_helper = plan_field(plan, kPlanHelper); + int broadcasts = 0; + + int trips = 0; + for (int guard = 0; guard < plan_ep * (plan_helpers + 2); ++guard) + { + ++trips; + bool const eligible = lane < plan_ep && !lane_frozen && lane_load > target; + int owner, owner_load; + if constexpr (NarrowKeys) + { + unsigned const key = __reduce_max_sync( + mask, eligible ? ((static_cast(lane_load) << 4) | (15u - static_cast(lane))) : 0u); + if (!key) + break; + owner = 15 - static_cast(key & 15u); + owner_load = static_cast(key >> 4); + } + else + { + owner_load = __reduce_max_sync(mask, eligible ? lane_load : -1); + if (owner_load < 0) + break; + owner = __reduce_min_sync(mask, eligible && lane_load == owner_load ? lane : INT_MAX); + } + + unsigned const packed + = __shfl_sync(mask, (static_cast(lane_send) << 4) | static_cast(lane_next), owner, 8); + int const owner_send = static_cast(packed >> 4); + int const owner_next = static_cast(packed & 15u); + if (owner_send >= plan_helpers || owner_next >= plan_helpers) + { + if (lane == owner) + lane_frozen = 1; + continue; + } + + /* owner and owner_next are now warp-uniform, so every lane can read + * candidates[] and histogram[] directly. Both are written by + * prepare_halo_m_sparse8 before its bar.sync and never inside this + * loop, so a uniform load returns exactly what the owner would have + * broadcast. The increment stays owner-only and stays ahead of the + * mass test, as before. */ + int const expert = candidates[owner * plan_helpers + owner_next]; + if (lane == owner) + ++lane_next; + int const mass = histogram[expert]; + if (mass == 0) + { + if (lane == owner) + lane_frozen = 1; + continue; + } + + int give = imin(owner_load - target, mass); + if (give <= 0) + { + if (lane == owner) + lane_frozen = 1; + continue; + } + + int positive_target = imax(target, 1); + int need = imax(1, mass / positive_target + (mass % positive_target != 0)); + int wanted = depth - 1; + for (int level = depth - 1; level >= 0; --level) + { + if (plan_ep / (1 << level) >= need) + { + wanted = level; + break; + } + } + + int chosen_level = -1; + int chosen_begin = -1; + int chosen_size = 0; + if constexpr (FixedEp == 8 && FixedHelpers == 4) + { + int slack = lane != owner ? imax(0, target - lane_load) : 0; + /* Two independent full-mask sums instead of one lane-varying-mask + * reduce followed by two broadcasts. Integer addition is exact, so + * masking the addend gives the same half sums; the two reduces do + * not depend on each other, so the chain is depth 1 instead of 3, + * and the redux.sync no longer runs with a per-lane membermask. */ + int room0 = __reduce_add_sync(mask, lane < 4 ? slack : 0); + int room4 = __reduce_add_sync(mask, lane >= 4 ? slack : 0); + unsigned receive_mask = __ballot_sync(mask, lane_recv < plan_helpers); + bool valid0 = (receive_mask & 0x0fu) == 0x0fu && room0 > 0; + bool valid4 = (receive_mask & 0xf0u) == 0xf0u && room4 > 0; + bool want_half = mass <= 4 * target; + if (want_half && (valid0 || valid4)) + { + chosen_level = 1; + chosen_begin = valid0 && (!valid4 || room0 >= room4) ? 0 : 4; + chosen_size = 4; + } + else if (receive_mask == mask && room0 + room4 > 0) + { + chosen_level = 0; + chosen_begin = 0; + chosen_size = 8; + } + } + else + { + for (int level = wanted; level >= 0 && chosen_level < 0; --level) + { + int size = plan_ep / (1 << level); + bool active = lane < plan_ep; + int begin = (lane / size) * size; + unsigned group_mask = rank_group_mask(begin, size); + int local_room = active && lane != owner ? imax(0, target - lane_load) : 0; + int room = 0; + if (active) + room = __reduce_add_sync(group_mask, local_room); + unsigned receive_mask = __ballot_sync(mask, active && lane_recv < plan_helpers); + bool group_leader = active && lane == begin && (receive_mask & group_mask) == group_mask; + unsigned leader_mask = __ballot_sync(mask, group_leader); + int best_room = -1; + int best_begin = INT_MAX; + if (leader_mask) + { + int const first_leader = __ffs(leader_mask) - 1; + if constexpr (NarrowKeys) + { + unsigned key = 0; + if (group_leader) + key = __reduce_max_sync( + leader_mask, (static_cast(room) << 4) | (15u - static_cast(begin))); + key = __shfl_sync(mask, key, first_leader, 8); + best_room = static_cast(key >> 4); + best_begin = 15 - static_cast(key & 15u); + } + else + { + if (group_leader) + { + best_room = __reduce_max_sync(leader_mask, room); + best_begin = __reduce_min_sync(leader_mask, room == best_room ? begin : INT_MAX); + } + best_room = __shfl_sync(mask, best_room, first_leader, 8); + best_begin = __shfl_sync(mask, best_begin, first_leader, 8); + } + } + if (best_begin != INT_MAX && best_room > 0) + { + chosen_level = level; + chosen_begin = best_begin; + chosen_size = size; + } + } + } + if (chosen_level < 0) + { + if (lane == owner) + lane_frozen = 1; + continue; + } + if (broadcasts >= kMaxBroadcasts) + break; + + int lane_quota = lane == owner ? mass : 0; + int remaining = give; + unsigned visited = 0; + for (int step = 0; step < chosen_size && remaining > 0; ++step) + { + bool can_receive = lane >= chosen_begin && lane < chosen_begin + chosen_size && lane != owner + && !(visited & (1u << lane)); + int destination, destination_load; + if constexpr (NarrowKeys) + { + unsigned const key = __reduce_min_sync(mask, + can_receive ? ((static_cast(lane_load) << 4) | static_cast(lane)) + : 0xffffffffu); + if (key == 0xffffffffu) + break; + destination = static_cast(key & 15u); + destination_load = static_cast(key >> 4); + } + else + { + destination_load = __reduce_min_sync(mask, can_receive ? lane_load : INT_MAX); + destination = __reduce_min_sync(mask, can_receive && lane_load == destination_load ? lane : INT_MAX); + if (destination == INT_MAX) + break; + } + visited |= 1u << destination; + int room = target - destination_load; + if (room <= 0) + break; + int amount = imin(room, remaining); + if (lane == owner) + lane_quota -= amount; + if (lane == destination) + lane_quota += amount; + remaining -= amount; + } + + int moved = give - remaining; + if (moved <= 0) + { + if (lane == owner) + lane_frozen = 1; + continue; + } + + int b = broadcasts++; + if (lane == 0) + { + p_expert[b] = expert; + p_level[b] = chosen_level; + p_begin[b] = chosen_begin; + p_size[b] = chosen_size; + p_owner[b] = owner; + p_helper[b] = -1; + } + if (lane < plan_ep) + quota[b * plan_ep + lane] = lane_quota; + + bool in_group = lane >= chosen_begin && lane < chosen_begin + chosen_size; + if (in_group) + ++lane_recv; + if (lane == owner) + { + ++lane_send; + lane_load -= moved; + } + else if (in_group) + { + lane_load += lane_quota; + } + } + return broadcasts; +} + +/* Select the key representation once, from the actual complete rank load. + * This guard keeps owner/destination loads and group-room keys representable. + * Keep the wide path for large valid int32 inputs without putting a runtime + * representation branch on each iteration of either planning loop. */ +template +__device__ __forceinline__ int plan_halo_m_register8( + int const* histogram, int const* load, int* quota, int const* candidates, int* plan, int ep, int helpers) +{ + int const plan_ep = FixedEp > 0 ? FixedEp : ep; + int const lane = threadIdx.x & 31; + int const lane_load = lane < plan_ep ? load[lane] : 0; + int const total = __reduce_add_sync(0xffu, lane_load); + if (total <= 0x0fffffff) + return plan_halo_m_register8_impl( + histogram, lane_load, total, quota, candidates, plan, ep, helpers); + return plan_halo_m_register8_impl( + histogram, lane_load, total, quota, candidates, plan, ep, helpers); +} + +/* Exact HALO-M state machine executed by warp 0. Each lane owns one rank; + * warp reductions replace all serial rank/group/destination scans. */ +__device__ int plan_halo_m_warp(int const* histogram, int* load, int* quota, int const* candidates, int* recv_used, + int* send_used, int* next_index, int* frozen, int* plan, int ep, int experts, int helpers) +{ + int const candidate_count = imin(helpers, experts / ep); + int const lane = threadIdx.x & 31; + int const reduce_width = ep <= 8 ? 8 : 32; + int const reduce_start = reduce_width >> 1; + unsigned const reduce_mask = ep <= 8 ? 0xffu : 0xffffffffu; + int total = lane < ep ? load[lane] : 0; + for (int offset = reduce_start; offset; offset >>= 1) + { + int const other = __shfl_down_sync(reduce_mask, total, offset, reduce_width); + if (lane + offset < reduce_width) + total += other; + } + total = __shfl_sync(reduce_mask, total, 0, reduce_width); + int const target = total / ep + (total % ep != 0); + int const depth = hierarchy_depth(ep); + int* p_expert = plan_field(plan, kPlanExpert); + int* p_level = plan_field(plan, kPlanLevel); + int* p_begin = plan_field(plan, kPlanGroupBegin); + int* p_size = plan_field(plan, kPlanGroupSize); + int* p_owner = plan_field(plan, kPlanOwner); + int* p_helper = plan_field(plan, kPlanHelper); + int broadcasts = 0; + + for (int guard = 0; guard < ep * (candidate_count + 2); ++guard) + { + bool eligible = lane < ep && !frozen[lane] && load[lane] > target; + int best_load = eligible ? load[lane] : -1; + int best_rank = eligible ? lane : INT_MAX; + for (int offset = reduce_start; offset; offset >>= 1) + { + int other_load = __shfl_down_sync(reduce_mask, best_load, offset, reduce_width); + int other_rank = __shfl_down_sync(reduce_mask, best_rank, offset, reduce_width); + if (lane + offset < reduce_width + && (other_load > best_load || (other_load == best_load && other_rank < best_rank))) + { + best_load = other_load; + best_rank = other_rank; + } + } + int owner = __shfl_sync(reduce_mask, best_rank, 0, reduce_width); + if (owner == INT_MAX) + break; + /* One lane owns the capacity check and shared cursor transition. + * Publish either the chosen expert or -1 for exhaustion explicitly; + * no reader can observe a partially advanced cursor. */ + int expert = -1; + if (lane == 0) + { + if (send_used[owner] >= helpers || next_index[owner] >= candidate_count) + { + frozen[owner] = 1; + } + else + { + int const candidate_index = next_index[owner]++; + expert = candidates[owner * candidate_count + candidate_index]; + } + } + expert = __shfl_sync(reduce_mask, expert, 0, reduce_width); + __syncwarp(reduce_mask); + if (expert < 0) + continue; + int mass = histogram[expert]; + if (mass == 0) + { + if (lane == 0) + frozen[owner] = 1; + __syncwarp(reduce_mask); + continue; + } + int give = imin(load[owner] - target, mass); + if (give <= 0) + { + if (lane == 0) + frozen[owner] = 1; + __syncwarp(reduce_mask); + continue; + } + + int need = imax(1, mass / imax(target, 1) + (mass % imax(target, 1) != 0)); + int wanted = depth - 1; + for (int level = depth - 1; level >= 0; --level) + { + if (ep / (1 << level) >= need) + { + wanted = level; + break; + } + } + int chosen_level = -1; + int chosen_begin = -1; + int chosen_size = 0; + for (int level = wanted; level >= 0 && chosen_level < 0; --level) + { + int size = ep / (1 << level); + int begin = lane * size; + bool active = begin < ep; + bool affordable = active; + int room = 0; + if (active) + { + for (int rank = begin; rank < begin + size; ++rank) + { + affordable &= recv_used[rank] < helpers; + if (rank != owner) + room += imax(0, target - load[rank]); + } + } + int candidate_room = affordable ? room : -1; + int candidate_begin = affordable ? begin : INT_MAX; + for (int offset = reduce_start; offset; offset >>= 1) + { + int other_room = __shfl_down_sync(reduce_mask, candidate_room, offset, reduce_width); + int other_begin = __shfl_down_sync(reduce_mask, candidate_begin, offset, reduce_width); + if (lane + offset < reduce_width + && (other_room > candidate_room || (other_room == candidate_room && other_begin < candidate_begin))) + { + candidate_room = other_room; + candidate_begin = other_begin; + } + } + int best_room = __shfl_sync(reduce_mask, candidate_room, 0, reduce_width); + int best_begin = __shfl_sync(reduce_mask, candidate_begin, 0, reduce_width); + if (best_begin != INT_MAX && best_room > 0) + { + chosen_level = level; + chosen_begin = best_begin; + chosen_size = size; + } + } + if (chosen_level < 0) + { + if (lane == 0) + frozen[owner] = 1; + __syncwarp(reduce_mask); + continue; + } + if (broadcasts >= kMaxBroadcasts) + break; + + int b = broadcasts++; + if (lane == 0) + { + p_expert[b] = expert; + p_level[b] = chosen_level; + p_begin[b] = chosen_begin; + p_size[b] = chosen_size; + p_owner[b] = owner; + p_helper[b] = -1; + } + if (lane < ep) + quota[b * ep + lane] = 0; + __syncwarp(reduce_mask); + if (lane == 0) + quota[b * ep + owner] = mass; + __syncwarp(reduce_mask); + + int remaining = give; + unsigned visited = 0; + for (int step = 0; step < chosen_size && remaining > 0; ++step) + { + bool can_receive = lane >= chosen_begin && lane < chosen_begin + chosen_size && lane != owner + && !(visited & (1u << lane)); + int cold_load = can_receive ? load[lane] : INT_MAX; + int cold_rank = can_receive ? lane : INT_MAX; + for (int offset = reduce_start; offset; offset >>= 1) + { + int other_load = __shfl_down_sync(reduce_mask, cold_load, offset, reduce_width); + int other_rank = __shfl_down_sync(reduce_mask, cold_rank, offset, reduce_width); + if (lane + offset < reduce_width + && (other_load < cold_load || (other_load == cold_load && other_rank < cold_rank))) + { + cold_load = other_load; + cold_rank = other_rank; + } + } + int destination = __shfl_sync(reduce_mask, cold_rank, 0, reduce_width); + if (destination == INT_MAX) + break; + visited |= 1u << destination; + int room = target - load[destination]; + if (room <= 0) + continue; + int amount = imin(room, remaining); + if (lane == 0) + { + quota[b * ep + owner] -= amount; + quota[b * ep + destination] += amount; + } + remaining -= amount; + __syncwarp(reduce_mask); + } + int moved = give - remaining; + if (moved <= 0) + { + --broadcasts; + if (lane == 0) + frozen[owner] = 1; + __syncwarp(reduce_mask); + continue; + } + if (lane >= chosen_begin && lane < chosen_begin + chosen_size) + ++recv_used[lane]; + if (lane == owner) + { + ++send_used[owner]; + load[owner] -= moved; + } + else if (lane >= chosen_begin && lane < chosen_begin + chosen_size) + { + load[lane] += quota[b * ep + lane]; + } + __syncwarp(reduce_mask); + } + return broadcasts; +} + +__device__ __forceinline__ int warp_sum(int value) +{ + for (int offset = 16; offset; offset >>= 1) + { + int const other = __shfl_down_sync(0xffffffffu, value, offset); + if ((threadIdx.x & 31) + offset < 32) + value += other; + } + return __shfl_sync(0xffffffffu, value, 0); +} + +__device__ __forceinline__ int warp_exclusive_scan(int value) +{ + int sum = value; + int const lane = threadIdx.x & 31; + for (int offset = 1; offset < 32; offset <<= 1) + { + int other = __shfl_up_sync(0xffffffffu, sum, offset); + if (lane >= offset) + sum += other; + } + return sum - value; +} + +__device__ void halo_q_repair(int receiver, int donor, int broadcasts, int const* order, int const* host_mask, + int const* counts, int const* p_expert, int* quota, int ep, int experts) +{ + int const lane = threadIdx.x & 31; + if (broadcasts > 32) + { + int local_pull = 0, local_give = 0; + for (int base = 0; base < broadcasts; base += 32) + { + int const b = base + lane < broadcasts ? order[base + lane] : -1; + if (b >= 0 && (host_mask[b] & (1u << receiver)) && (host_mask[b] & (1u << donor))) + { + int const expert = p_expert[b]; + int const cr = counts[receiver * experts + expert]; + int const cd = counts[donor * experts + expert]; + int const qr = quota[b * ep + receiver]; + int const qd = quota[b * ep + donor]; + local_pull += imin(imax(0, cr - qr), imax(0, qd - cd)); + local_give += imax(0, qr - cr); + } + } + int const delta = imin(warp_sum(local_pull), warp_sum(local_give)); + int pull_before_tiles = 0, give_before_tiles = 0; + for (int base = 0; base < broadcasts; base += 32) + { + int const b = base + lane < broadcasts ? order[base + lane] : -1; + int pull = 0, give = 0; + if (b >= 0 && (host_mask[b] & (1u << receiver)) && (host_mask[b] & (1u << donor))) + { + int const expert = p_expert[b]; + int const cr = counts[receiver * experts + expert]; + int const cd = counts[donor * experts + expert]; + int const qr = quota[b * ep + receiver]; + int const qd = quota[b * ep + donor]; + pull = imin(imax(0, cr - qr), imax(0, qd - cd)); + give = imax(0, qr - cr); + } + int const pull_before = pull_before_tiles + warp_exclusive_scan(pull); + int const give_before = give_before_tiles + warp_exclusive_scan(give); + int const take = imin(pull, imax(0, delta - pull_before)); + int const release = imin(give, imax(0, delta - give_before)); + pull_before_tiles += warp_sum(pull); + give_before_tiles += warp_sum(give); + if (b >= 0) + { + quota[b * ep + receiver] += take - release; + quota[b * ep + donor] += release - take; + } + __syncwarp(); + } + return; + } + int b = lane < broadcasts ? order[lane] : -1; + int pull = 0; + int give = 0; + if (b >= 0 && (host_mask[b] & (1u << receiver)) && (host_mask[b] & (1u << donor))) + { + int expert = p_expert[b]; + int cr = counts[receiver * experts + expert]; + int cd = counts[donor * experts + expert]; + int qr = quota[b * ep + receiver]; + int qd = quota[b * ep + donor]; + pull = imin(imax(0, cr - qr), imax(0, qd - cd)); + give = imax(0, qr - cr); + } + int delta = imin(warp_sum(pull), warp_sum(give)); + int pull_before = warp_exclusive_scan(pull); + int give_before = warp_exclusive_scan(give); + int pull_take = imin(pull, imax(0, delta - pull_before)); + int give_take = imin(give, imax(0, delta - give_before)); + if (b >= 0) + { + quota[b * ep + receiver] += pull_take - give_take; + quota[b * ep + donor] += give_take - pull_take; + } + __syncwarp(); +} + +__device__ __forceinline__ int warp_sum8(int value) +{ + return __reduce_add_sync(0xffu, value); +} + +__device__ __forceinline__ int warp_exclusive_scan8(int value) +{ + int sum = value; + int const lane = threadIdx.x & 7; + for (int offset = 1; offset < 8; offset <<= 1) + { + int other = __shfl_up_sync(0xffu, sum, offset, 8); + if (lane >= offset) + sum += other; + } + return sum - value; +} + +__device__ __forceinline__ void halo_q_direction8( + bool common, int receiver_count, int donor_count, int& receiver_quota, int& donor_quota) +{ + int pull = common ? imin(imax(0, receiver_count - receiver_quota), imax(0, donor_quota - donor_count)) : 0; + int give = common ? imax(0, receiver_quota - receiver_count) : 0; + int pull_total = warp_sum8(pull); + int give_total = warp_sum8(give); + int delta = imin(pull_total, give_total); + if (!delta) + return; + int pull_before = warp_exclusive_scan8(pull); + int give_before = warp_exclusive_scan8(give); + int pull_take = imin(pull, imax(0, delta - pull_before)); + int give_take = imin(give, imax(0, delta - give_before)); + receiver_quota += pull_take - give_take; + donor_quota += give_take - pull_take; +} + +/* Fuse both ordered directions of one rank pair. Each broadcast lane loads + * its two count/quota entries once, executes left<-right then right<-left in + * registers, and writes the final pair once. + * + * `b`, `hosts` and `expert` are the lane's own broadcast identity and do not + * change between phases or rounds, so the caller hoists them. Reading them + * here put a four-deep dependent load chain (order[lane] -> b -> host_mask[b] + * and the GLOBAL p_expert[b] -> expert -> counts[...]) at the head of all + * fourteen phases. */ +__device__ __forceinline__ void halo_q_pair8( + int left, int right, int b, int hosts, int expert, int const* counts, int* quota, int ep, int experts) +{ + int const lane = threadIdx.x & 31; + if (lane >= 8) + return; + bool common = false; + int count_left = 0; + int count_right = 0; + int quota_left = 0; + int quota_right = 0; + if (b >= 0) + { + common = (hosts & (1u << left)) && (hosts & (1u << right)); + if (common) + { + count_left = counts[left * experts + expert]; + count_right = counts[right * experts + expert]; + quota_left = quota[b * ep + left]; + quota_right = quota[b * ep + right]; + } + } + + halo_q_direction8(common, count_left, count_right, quota_left, quota_right); + halo_q_direction8(common, count_right, count_left, quota_right, quota_left); + + if (common) + { + quota[b * ep + left] = quota_left; + quota[b * ep + right] = quota_right; + } +} + +__device__ __forceinline__ void halo_q_active_barrier(int threads) +{ + asm volatile("bar.sync 1, %0;" ::"r"(threads) : "memory"); +} + +__device__ __forceinline__ void scheduler_tail_barrier() +{ + asm volatile("bar.sync 3, 64;" ::: "memory"); +} + +/* Build order[] and host_mask[] with one lane per broadcast. + * + * This replaces serial insertion-sort and mask-construction loops while + * preserving the same stable order. + * + * order[]: lane b's position is the number of broadcasts that sort before it. + * The `j < lane` tiebreak makes this the same stable order the + * insertion sort produced (its shift test is a strict `>`), so the + * result is bit-identical, not merely equivalent. + * host_mask: the destination group is a contiguous rank range, so the mask is + * closed form; no loop over group members is needed. + * rank_group_mask handles the full EP32 word without a shift by32. + */ +__device__ __forceinline__ void halo_q_build_order_and_masks(int const* p_expert, int const* p_begin, int const* p_size, + int const* p_owner, int* order, int* host_mask, int broadcasts) +{ + int const lane = threadIdx.x & 31; + if (broadcasts > 32) + { + for (int b = lane; b < broadcasts; b += 32) + { + int const mine = p_expert[b]; + int position = 0; + for (int j = 0; j < broadcasts; ++j) + position += p_expert[j] < mine || (p_expert[j] == mine && j < b); + order[position] = b; + host_mask[b] = static_cast(rank_group_mask(p_begin[b], p_size[b]) | (1u << p_owner[b])); + } + return; + } + /* Read p_expert cooperatively and distribute values with shuffles. All + * lanes must participate before inactive lanes return; inactive lanes + * carry INT_MAX and never contribute to a real comparison. */ + bool const active = lane < broadcasts; + int const mine = active ? p_expert[lane] : INT_MAX; + int position = 0; + for (int j = 0; j < broadcasts; ++j) + { + int const other = __shfl_sync(0xffffffffu, mine, j); + position += (other < mine) || (other == mine && j < lane); + } + if (!active) + return; + order[position] = lane; + host_mask[lane] = static_cast(rank_group_mask(p_begin[lane], p_size[lane]) | (1u << p_owner[lane])); +} + +template +__device__ void run_halo_q(int const* counts, int* quota, int* plan, int broadcasts, int ep, int experts, int* ring, + int* order, int* host_mask) +{ + int const plan_ep = FixedEp > 0 ? FixedEp : ep; + int const tid = threadIdx.x; + int const warp = tid >> 5; + int const lane = tid & 31; + int* p_expert = plan_field(plan, kPlanExpert); + int* p_begin = plan_field(plan, kPlanGroupBegin); + int* p_size = plan_field(plan, kPlanGroupSize); + int* p_owner = plan_field(plan, kPlanOwner); + if (plan_ep <= 8 && broadcasts <= 8) + { + if (warp < plan_ep / 2) + { + if (warp == 0) + halo_q_build_order_and_masks(p_expert, p_begin, p_size, p_owner, order, host_mask, broadcasts); + halo_q_active_barrier((plan_ep / 2) * 32); + /* Loop-invariant broadcast identity, read once after the barrier + * that publishes order[] and host_mask[]. All four participating + * warps read the same lane-indexed entries. */ + int const my_b = lane < broadcasts ? order[lane] : -1; + int const my_hosts = my_b >= 0 ? host_mask[my_b] : 0; + int const my_expert = my_b >= 0 ? p_expert[my_b] : -1; + if (broadcasts >= 2) + { + int const ring_size = plan_ep - 1; + /* Use two unconditional repair rounds. Phases with no required repair + * are no-ops, so this preserves the conditional algorithm's result. */ + int const repair_rounds = 2; + for (int round = 0; round < repair_rounds; ++round) + { + for (int phase = 0; phase < plan_ep - 1; ++phase) + { + int right_index = warp == 0 ? 0 : plan_ep - 1 - warp; + int right_position = right_index - phase; + if (right_position < 0) + right_position += ring_size; + int right = 1 + right_position; + int left = 0; + if (warp != 0) + { + int left_position = warp - phase; + if (left_position < 0) + left_position += ring_size; + left = 1 + left_position; + } + halo_q_pair8(left, right, my_b, my_hosts, my_expert, counts, quota, plan_ep, experts); + if constexpr (FixedEp == 8) + { + /* Must stay unconditional. Omitting it after a + * round's last phase races twice: warp 0 enters the + * next round's phase 0 while warp 3 still writes the + * same quota[], and after the final phase warps 1-3 + * leave run_halo_q while warp 0 -- which only meets + * warp 8 at scheduler_tail_barrier -- walks into + * build_route_prefix and reads every quota[]. */ + halo_q_active_barrier(128); + } + else + { + halo_q_active_barrier((plan_ep / 2) * 32); + } + } + } + } + } + /* The unconditional repair rounds above replace the conditional scan. */ + return; + } + if (warp == 0) + halo_q_build_order_and_masks(p_expert, p_begin, p_size, p_owner, order, host_mask, broadcasts); + if (tid < plan_ep - 1) + ring[tid] = tid + 1; + __syncthreads(); + if (broadcasts < 2) + return; + + for (int round = 0; round < 2; ++round) + { + for (int phase = 0; phase < plan_ep - 1; ++phase) + { + if (warp < plan_ep / 2) + { + int left, right; + if (warp == 0) + { + left = 0; + right = ring[0]; + } + else + { + left = ring[warp]; + right = ring[plan_ep - 1 - warp]; + } + halo_q_repair(left, right, broadcasts, order, host_mask, counts, p_expert, quota, plan_ep, experts); + halo_q_repair(right, left, broadcasts, order, host_mask, counts, p_expert, quota, plan_ep, experts); + } + __syncthreads(); + if (tid == 0) + { + int last = ring[plan_ep - 2]; + for (int index = plan_ep - 2; index > 0; --index) + ring[index] = ring[index - 1]; + ring[0] = last; + } + __syncthreads(); + } + } + (void) lane; +} + +__device__ void color_and_publish_local(int* plan, int broadcasts, int ep, int helpers, int local_rank, int* used, + int* out_ids, int* out_levels, int* out_owners, int* status, Epoch epoch) +{ + int* p_expert = plan_field(plan, kPlanExpert); + int* p_level = plan_field(plan, kPlanLevel); + int* p_begin = plan_field(plan, kPlanGroupBegin); + int* p_size = plan_field(plan, kPlanGroupSize); + int* p_owner = plan_field(plan, kPlanOwner); + int* p_helper = plan_field(plan, kPlanHelper); + int const active_helpers = imin(helpers, kMaxBroadcasts); + int const words = (active_helpers + 31) / 32; + for (int rank = 0; rank < ep; ++rank) + for (int word = 0; word < words; ++word) + used[rank * kHelperMaskWords + word] = 0; + for (int level = 0; level < hierarchy_depth(ep); ++level) + { + for (int b = 0; b < broadcasts; ++b) + { + if (p_level[b] != level) + continue; + unsigned unavailable[kHelperMaskWords] = {0}; + for (int rank = p_begin[b]; rank < p_begin[b] + p_size[b]; ++rank) + for (int word = 0; word < words; ++word) + unavailable[word] |= static_cast(used[rank * kHelperMaskWords + word]); + int helper = -1; + for (int slot = 0; slot < active_helpers; ++slot) + if (!(unavailable[slot / 32] & (1u << (slot % 32)))) + { + helper = slot; + break; + } + if (helper < 0) + { + record_status(status + 4, epoch); + helper = 0; + } + p_helper[b] = helper; + for (int rank = p_begin[b]; rank < p_begin[b] + p_size[b]; ++rank) + used[rank * kHelperMaskWords + helper / 32] |= static_cast(1u << (helper % 32)); + } + } + for (int slot = 0; slot < helpers; ++slot) + { + out_ids[slot] = -1; + out_levels[slot] = -1; + out_owners[slot] = -1; + } + for (int b = 0; b < broadcasts; ++b) + { + if (!group_contains(local_rank, p_begin[b], p_size[b])) + continue; + int helper = p_helper[b]; + if (out_ids[helper] != -1) + record_status(status + 4, epoch); + out_ids[helper] = p_expert[b]; + out_levels[helper] = p_level[b]; + out_owners[helper] = p_owner[b]; + } +} + +__device__ void build_route_prefix(int const* counts, int const* quota, int* prefix, int const* plan, int broadcasts, + int ep, int experts, int local_rank, int participating_threads) +{ + int* p_expert = plan_field(const_cast(plan), kPlanExpert); + for (int b = threadIdx.x; b < broadcasts; b += participating_threads) + { + int expert = p_expert[b]; + int source_before = 0; + for (int rank = 0; rank < local_rank; ++rank) + { + int diagonal = imin(counts[rank * experts + expert], quota[b * ep + rank]); + source_before += counts[rank * experts + expert] - diagonal; + } + int source_count = counts[local_rank * experts + expert]; + int source_diagonal = imin(source_count, quota[b * ep + local_rank]); + int residual_supply = source_count - source_diagonal; + /* The loop used to run to destination == ep. route_destination only + * ever loads prefix[destination + step] under `candidate < ep`, so the + * ep entry is never read, and the accumulation at ep-1 fed nothing but + * that dead entry. Both are gone; entry 0 is still written because it + * costs one store and leaving it stale would depend on nobody ever + * reading it, which is a weaker guarantee than writing the right value. */ + int destination_before = 0; + for (int destination = 0; destination < ep; ++destination) + { + int overlap = destination_before - source_before; + overlap = imax(0, imin(residual_supply, overlap)); + int diagonal_prefix = local_rank < destination ? source_diagonal : 0; + prefix[b * (ep + 1) + destination] = diagonal_prefix + overlap; + if (destination + 1 < ep) + { + int diagonal = imin(counts[destination * experts + expert], quota[b * ep + destination]); + destination_before += quota[b * ep + destination] - diagonal; + } + } + } +} + +__device__ __forceinline__ int* worker_expert_lut(int experts) +{ + extern __shared__ int shared[]; + return shared + (blockDim.x >> 5) * (experts + 1); +} + +__device__ void prepare_worker_default_slots(int* expert_lut, int ep, int experts, int helpers) +{ + int const home = experts / ep; + int const local_slots = home + helpers; + for (int expert = threadIdx.x; expert < experts; expert += blockDim.x) + { + int owner = expert / home; + int local = expert - owner * home; + /* Negative values distinguish the precomputed cold slot from a + * nonnegative broadcast index. Slot zero therefore maps to -1. */ + expert_lut[expert] = -(owner * local_slots + local) - 1; + } + __syncthreads(); +} + +__device__ void prepare_worker_selected_experts(int* expert_lut, int const* plan) +{ + int const broadcasts = plan[1]; + int const* p_expert = plan + 2 + kPlanExpert * kMaxBroadcasts; + for (int b = threadIdx.x; b < broadcasts; b += blockDim.x) + expert_lut[p_expert[b]] = b; + __syncthreads(); +} + +__device__ bool wait_for_plan_epoch(int* marker, Epoch epoch, unsigned long long spin_cycles, int* status, int* ready) +{ + if (threadIdx.x == 0) + { + *ready = 1; + unsigned long long start = clock64(); + Epoch observed = load_relaxed_device(marker); + while (epoch_before(observed, epoch)) + { + if (clock64() - start > spin_cycles) + { + record_status(status + 2, epoch); + *ready = 0; + break; + } + observed = load_relaxed_device(marker); + } + if (!epoch_before(observed, epoch)) + (void) load_acquire_device(marker); + } + __syncthreads(); + return *ready != 0; +} + +/* Prefix entries are monotonic. Binary lifting returns the same number of + * thresholds <= ordinal as the old linear scan, including repeated entries. */ +__device__ __forceinline__ int route_destination(int const* prefix, int ordinal, int ep) +{ + int destination = 0; +#pragma unroll + for (int step = 16; step; step >>= 1) + { + int candidate = destination + step; + if (candidate < ep && ordinal >= prefix[candidate]) + destination = candidate; + } + return destination; +} + +/* Once HALO-M publishes its selected experts, workers can finish every cold + * route while CTA0 runs HALO-Q. Selected route indices are compacted into + * disjoint slices of route_aux; their ordinal values remain in output. */ +__device__ void pre_materialize_and_compact_worker_routes(int const* routes, int* output, int* route_aux, int* status, + int route_count, int experts, int worker, int workers, Epoch epoch) +{ + extern __shared__ int shared[]; + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const warps = blockDim.x >> 5; + int const bins = experts + 1; + int* expert_lut = worker_expert_lut(experts); + int* hot_indices = route_aux + 2 + workers * bins; + + int const worker_span = (route_count + workers - 1) / workers; + int const worker_begin = worker * worker_span; + int const worker_end = imin(route_count, worker_begin + worker_span); + int const warp_span = ((worker_span + warps * 32 - 1) / (warps * 32)) * 32; + int const begin = worker_begin + warp * warp_span; + int const end = imin(worker_end, begin + warp_span); + int hot_count = 0; + for (int base = begin; base < end; base += 32) + { + int index = base + lane; + bool active = index < end; + int expert = active ? routes[index] : -1; + int slot = -1; + bool hot = false; + if (active && expert >= 0 && expert < experts) + { + int entry = expert_lut[expert]; + hot = entry >= 0; + if (!hot) + slot = -entry - 1; + } + else if (active && expert != -1) + { + record_status(status + 3, epoch); + } + unsigned hot_mask = __ballot_sync(0xffffffffu, hot); + if (hot) + { + int offset = __popc(hot_mask & ((1u << lane) - 1u)); + hot_indices[begin + hot_count + offset] = index; + } + else if (active) + { + output[index] = slot; + } + hot_count += __popc(hot_mask); + } + if (lane == 0) + shared[warp * bins + experts] = hot_count; +} + +/* The final epoch only has to resolve selected routes. Stable ordinal + * cursors remain in the per-warp shared histogram from the initial pass. */ +__device__ void materialize_worker_hot_routes(int const* routes, int* output, int const* route_prefix, int const* plan, + int* route_aux, int* status, int route_count, int ep, int experts, int helpers, int worker, int workers, + Epoch epoch) +{ + extern __shared__ int shared[]; + int const lane = threadIdx.x & 31; + int const warp = threadIdx.x >> 5; + int const warps = blockDim.x >> 5; + int const bins = experts + 1; + int* expert_lut = worker_expert_lut(experts); + int* hot_indices = route_aux + 2 + workers * bins; + int const home = experts / ep; + int const local_slots = home + helpers; + int const* p_begin = plan + 2 + kPlanGroupBegin * kMaxBroadcasts; + int const* p_size = plan + 2 + kPlanGroupSize * kMaxBroadcasts; + int const* p_owner = plan + 2 + kPlanOwner * kMaxBroadcasts; + int const* p_helper = plan + 2 + kPlanHelper * kMaxBroadcasts; + + int const worker_span = (route_count + workers - 1) / workers; + int const worker_begin = worker * worker_span; + int const worker_end = imin(route_count, worker_begin + worker_span); + int const warp_span = ((worker_span + warps * 32 - 1) / (warps * 32)) * 32; + int const begin = worker_begin + warp * warp_span; + int const end = imin(worker_end, begin + warp_span); + int hot_count = shared[warp * bins + experts]; + hot_count = imin(hot_count, imax(0, end - begin)); + for (int offset = lane; offset < hot_count; offset += 32) + { + int index = hot_indices[begin + offset]; + int expert = routes[index]; + int b = expert_lut[expert]; + int ordinal = output[index] + shared[warp * bins + expert]; + int const* prefix = route_prefix + b * (ep + 1); + int destination = route_destination(prefix, ordinal, ep); + int slot = -1; + if (group_contains(destination, p_begin[b], p_size[b])) + { + slot = destination * local_slots + home + p_helper[b]; + } + else if (destination == p_owner[b]) + { + slot = destination * local_slots + expert - p_owner[b] * home; + } + else + { + record_status(status + 5, epoch); + } + output[index] = slot; + } +} + +__device__ void materialize_worker_routes(int const* routes, int* output, int const* route_prefix, int const* plan, + int* status, int route_count, int ep, int experts, int helpers, int worker, int workers, Epoch epoch) +{ + extern __shared__ int shared[]; + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const warps = blockDim.x >> 5; + int const bins = experts + 1; + int* expert_lut = worker_expert_lut(experts); + int const home = experts / ep; + int const local_slots = home + helpers; + int const* p_begin = plan + 2 + kPlanGroupBegin * kMaxBroadcasts; + int const* p_size = plan + 2 + kPlanGroupSize * kMaxBroadcasts; + int const* p_owner = plan + 2 + kPlanOwner * kMaxBroadcasts; + int const* p_helper = plan + 2 + kPlanHelper * kMaxBroadcasts; + + int const worker_span = (route_count + workers - 1) / workers; + int const worker_begin = worker * worker_span; + int const worker_end = imin(route_count, worker_begin + worker_span); + int const warp_span = ((worker_span + warps * 32 - 1) / (warps * 32)) * 32; + int const begin = worker_begin + warp * warp_span; + int const end = imin(worker_end, begin + warp_span); + for (int index = begin + lane; index < end; index += 32) + { + int expert = routes[index]; + int slot = -1; + if (expert >= 0 && expert < experts) + { + int entry = expert_lut[expert]; + if (entry < 0) + { + slot = -entry - 1; + } + else + { + int b = entry; + int ordinal = output[index] + shared[warp * bins + expert]; + int const* prefix = route_prefix + b * (ep + 1); + int destination = route_destination(prefix, ordinal, ep); + if (group_contains(destination, p_begin[b], p_size[b])) + { + slot = destination * local_slots + home + p_helper[b]; + } + else if (destination == p_owner[b]) + { + slot = destination * local_slots + expert - p_owner[b] * home; + } + else + { + record_status(status + 5, epoch); + } + } + } + else if (expert != -1) + { + record_status(status + 3, epoch); + } + output[index] = slot; + } +} + +/* One scheduler generation on a grid of ctas resident CTAs of 512 threads: + * CTA 0 plans and publishes the helper placement, the other CTAs write routes. + * Callers launch it as halo_q_scheduler_kernel or as the first phase of the + * fused scheduler + in-switch copy kernel. */ +template +__device__ __forceinline__ void halo_q_scheduler_body(int const* routes, int* out_slots, int* out_ids, int* out_levels, + int* out_owners, unsigned long long const* peer_bases, int* status, int* partial, int* route_aux, int* grid_sync, + int* plan, int* route_prefix, int* plan_channel, int runtime_ep, int runtime_experts, int runtime_helpers, + int local_rank, int route_capacity, int ctas, int enable_pdl, unsigned long long spin_cycles, int plan_abi_version, + int plan_channel_words, int route_features, int capacity, int route_count) +{ + int const ep = FixedEp > 0 ? FixedEp : runtime_ep; + int const experts = FixedExperts > 0 ? FixedExperts : runtime_experts; + int const helpers = FixedHelpers > 0 ? FixedHelpers : runtime_helpers; + (void) plan_channel_words; // Validated by the host launch wrapper. + /* The scheduler grid is deliberately residency-safe (one CTA per SM and + * CTA count <= SM count), so signal as soon as every CTA starts. The + * independent same-stream quantizer may then occupy all unused SMs while + * these CTAs retain their own SMs for the complete scheduler lifetime. */ + if (enable_pdl && threadIdx.x == 0) + cudaTriggerProgrammaticLaunchCompletion(); + // Only the caller's valid prefix is readable. Keep the public output at + // its configured capacity, clearing stale routes inside this same kernel. + // Empty ranks still execute every collective; no route load occurs at N=0. + for (int index = route_count + blockIdx.x * blockDim.x + threadIdx.x; index < route_capacity; + index += gridDim.x * blockDim.x) + out_slots[index] = -1; + local_histogram_pass(routes, partial, route_count, experts, ctas); + Epoch epoch; + if (!grid_rendezvous(grid_sync, ctas, spin_cycles, status, &epoch)) + return; + + if (blockIdx.x == 0) + { + extern __shared__ int shared[]; + int offset = ep * experts; + int* counts = shared; + int* histogram = shared + offset; + offset += experts; + int* load = shared + offset; + offset += ep; + // Reuse the exact capacity already validated/computed by launch(). + int* quota = shared + offset; + offset += capacity * ep; + int* selected_flags = shared + offset; + offset += experts; + int* recv_used = shared + offset; + offset += ep * kHelperMaskWords; + int* send_used = shared + offset; + offset += ep; + int* next_index = shared + offset; + offset += ep; + int* frozen = shared + offset; + offset += ep; + int* ring = shared + offset; + offset += ep; + int* order = shared + offset; + offset += capacity; + int* host_mask = shared + offset; + + bool exchange_ok; + if (ep == 8) + exchange_ok = exchange_counts<8>( + peer_bases, partial, shared, ep, experts, local_rank, ctas, epoch, spin_cycles, status); + else + exchange_ok = exchange_counts<0>( + peer_bases, partial, shared, ep, experts, local_rank, ctas, epoch, spin_cycles, status); + if (!exchange_ok) + return; + /* EP8 keeps only its candidate masses; other EPs retain the dense + * histogram required by the generic planner path. */ + if (ep == 8 && helpers == 4 && experts >= 32) + { + prepare_halo_m_sparse8(counts, histogram, load, selected_flags, experts); + } + else + { + prepare_halo_m<0, 0>(counts, histogram, load, selected_flags, recv_used, send_used, next_index, frozen, ep, + experts, helpers); + } + + int broadcasts = 0; + if (threadIdx.x < (ep <= 8 ? 8 : 32)) + { + // These short-circuit bounds keep helpers*ep <= 120. + if (ep <= 8 && helpers <= 15 && helpers * ep <= experts) + { + if (ep == 8 && helpers == 4) + { + broadcasts = plan_halo_m_register8<8, 4>(histogram, load, quota, selected_flags, plan, ep, helpers); + } + else if (ep == 4 && helpers == 4) + { + broadcasts = plan_halo_m_register8<4, 4>(histogram, load, quota, selected_flags, plan, ep, helpers); + } + else + { + broadcasts = plan_halo_m_register8<0, 0>(histogram, load, quota, selected_flags, plan, ep, helpers); + } + } + else + { + broadcasts = plan_halo_m_warp(histogram, load, quota, selected_flags, recv_used, send_used, next_index, + frozen, plan, ep, experts, helpers); + } + if (threadIdx.x == 0) + plan[1] = broadcasts; + } + __syncthreads(); + broadcasts = plan[1]; + /* p_expert and plan[1] are produced by thread zero. Publish this + * placement-only epoch before HALO-Q so workers can construct their + * selected-expert LUT while quota repair is still running. */ + if (ctas > 1 && threadIdx.x == 0) + { + store_release_device(plan + 0, epoch); + } + /* Small-EP repair excludes warp8 and can overlap coloring. Generic + * repair first meets the coloring warp at its full-block barrier. */ + if (threadIdx.x == 8 * 32) + color_and_publish_local( + plan, broadcasts, ep, helpers, local_rank, recv_used, out_ids, out_levels, out_owners, status, epoch); + if (ep == 8) + run_halo_q<8>(counts, quota, plan, broadcasts, ep, experts, ring, order, host_mask); + else + run_halo_q<0>(counts, quota, plan, broadcasts, ep, experts, ring, order, host_mask); + if (ctas > 1) + { + int const warp = threadIdx.x >> 5; + /* build_route_prefix is independent of warp 8's outputs. Defer their + * rendezvous until both paths reach their first shared consumer so + * prefix construction and coloring can overlap safely. */ + if (warp == 0) + { + build_route_prefix(counts, quota, route_prefix, plan, broadcasts, ep, experts, local_rank, 32); + __syncwarp(); + } + if (warp == 0 || warp == 8) + scheduler_tail_barrier(); + if (warp == 0) + { + /* Release the worker CTAs here rather than after the channel + * publish. They gate on grid_sync+2 and consume route_prefix, + * which is complete above; plan_channel is consumed only by the + * copy stack, which acquires on the channel's own epoch word + * written with store_release_system below. Publishing after + * the release lets the workers coloring pass -- the work that + * actually determines when the kernel ends -- overlap the whole + * publish instead of waiting behind it. */ + if (threadIdx.x == 0) + store_release_device(grid_sync + 2, epoch); + sami_copy::publish_plan_channel_warp(plan_channel, out_ids, out_levels, out_owners, + copy_placement_view(plan), helpers, ep, local_rank, experts, plan_abi_version, route_features); + __syncwarp(); + } + return; + } + if (threadIdx.x == 0) + store_release_device(plan + 0, epoch); + __syncthreads(); + build_route_prefix(counts, quota, route_prefix, plan, broadcasts, ep, experts, local_rank, blockDim.x); + __syncthreads(); + sami_copy::publish_plan_channel_warp(plan_channel, out_ids, out_levels, out_owners, copy_placement_view(plan), + helpers, ep, local_rank, experts, plan_abi_version, route_features); + /* The block barrier orders route-prefix/helper writes before release. */ + __syncthreads(); + } + else + { + int const worker = blockIdx.x - 1; + int const workers = ctas - 1; + if (!stable_worker_ordinal_pass( + routes, out_slots, route_aux, status, route_count, experts, worker, workers, spin_cycles)) + return; + int* expert_lut = worker_expert_lut(experts); + prepare_worker_default_slots(expert_lut, ep, experts, helpers); + __shared__ int plan_ready; + if (!wait_for_plan_epoch(plan + 0, epoch, spin_cycles, status, &plan_ready)) + return; + prepare_worker_selected_experts(expert_lut, plan); + pre_materialize_and_compact_worker_routes( + routes, out_slots, route_aux, status, route_count, experts, worker, workers, epoch); + if (!wait_for_plan_epoch(grid_sync + 2, epoch, spin_cycles, status, &plan_ready)) + return; + materialize_worker_hot_routes(routes, out_slots, route_prefix, plan, route_aux, status, route_count, ep, + experts, helpers, worker, workers, epoch); + return; + } + + /* CTA-count one remains a correct serial fallback. */ + if (!stable_worker_ordinal_pass(routes, out_slots, route_aux, status, route_count, experts, 0, 1, spin_cycles)) + return; + int* expert_lut = worker_expert_lut(experts); + prepare_worker_default_slots(expert_lut, ep, experts, helpers); + prepare_worker_selected_experts(expert_lut, plan); + materialize_worker_routes( + routes, out_slots, route_prefix, plan, status, route_count, ep, experts, helpers, 0, 1, epoch); +} + +/* Dynamic shared-memory ints that the planner CTA and a worker CTA need. */ +__host__ __device__ __forceinline__ int halo_q_shared_ints(int ep, int experts, int capacity, int threads) +{ + int const workerInts = threads / 32 * (experts + 1) + experts; + int const plannerInts + = ep * experts + experts + ep + capacity * ep + experts + (kHelperMaskWords + 4) * ep + 2 * capacity; + return workerInts > plannerInts ? workerInts : plannerInts; +} + +} // namespace + +} // namespace kernels + +TRTLLM_NAMESPACE_END + +#endif // TRTLLM_KERNELS_MOE_REBALANCE_HALO_Q_KERNEL_CUH diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlaceholder.cpp b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlaceholder.cpp index dd1b5f075985..40ca0ed5ae89 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlaceholder.cpp +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlaceholder.cpp @@ -28,6 +28,11 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { +void checkMoeRebalanceHaloQParams(MoeRebalanceHaloQParams const&) +{ + TLLM_THROW("Per-layer dynamic EPLB HALO-Q requires a requested SM100, SM103, or SM107 target"); +} + void invokeMoeRebalanceHaloQ(MoeRebalanceHaloQParams const&, cudaStream_t) { TLLM_THROW("Per-layer dynamic EPLB HALO-Q requires a requested SM100, SM103, or SM107 target"); @@ -69,7 +74,8 @@ extern "C" int megamoe_tma_copy_bind_gpu_direct( return unsupported(); } -extern "C" int megamoe_tma_copy_submit_gpu_direct(MegamoeTmaCopyState*, uint64_t, uint64_t, void*) +int megamoe_tma_copy_submit_fused(MegamoeTmaCopyState*, tensorrt_llm::kernels::MoeRebalanceHaloQParams const&, + uint64_t*, uint64_t, uint64_t, uint64_t, void*) { return unsupported(); } diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu index bef362a9bdc6..1d5d4dfdd418 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu @@ -24,7 +24,12 @@ * Per-bank alternating tickets protect slot metadata and shared-memory reuse. * This permits every allocated slot to participate even when one warp owns * many slots, without exceeding the ISA's outstanding bulk-group limit. + * + * The copy runs as the second phase of one kernel whose first phase is the + * HALO-Q scheduler that produces its plan. */ +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h" +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh" #include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlanPolicy.h" #include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h" #include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlanBuilder.cuh" @@ -492,12 +497,14 @@ __device__ __forceinline__ void tma_copy_body(MegamoeTmaCopySegment const* segme int const physical_warp = threadIdx.x / 32, lane = threadIdx.x % 32; int const warp = physical_warp / 2; // logical RR worker, shared by its two warps bool const producer = (physical_warp & 1) == 0; + // Warps beyond the 2 * warps copy workers only join the CTA barriers. + bool const worker = warp < warps; int const slots_per_warp = megamoe_tma_warp_slot_count(total_slots, warps, warp); int const first_slot = megamoe_tma_warp_slot_begin(total_slots, warps, warp); SlotControl* control = reinterpret_cast(smem + static_cast(total_slots) * kSlotDataBytes); SlotControl* warp_control = control + first_slot; unsigned char* warp_buffer = smem + static_cast(first_slot) * kSlotDataBytes; - if (producer && lane == 0) + if (worker && producer && lane == 0) { for (int slot = 0; slot < slots_per_warp; ++slot) initialize_barrier(warp_control + slot); @@ -505,7 +512,7 @@ __device__ __forceinline__ void tma_copy_body(MegamoeTmaCopySegment const* segme } __syncthreads(); // Elect once for the entire pipeline so ptxas can prove one TMA issuer per warp. - if (cuda::ptx::elect_sync(0xffffffffu)) + if (worker && cuda::ptx::elect_sync(0xffffffffu)) { // CTA/SM is the fastest-varying dimension, then warp, then sequence. const uint64_t ordinal_base = megamoe_tma_worker_first_slice(blockIdx.x, gridDim.x, warp); @@ -601,7 +608,7 @@ struct GpuDirectStorage MegamoeTmaGpuPlanResult* results; }; -__global__ void tma_copy_gpu_direct_kernel(GpuDirectStorage direct, int warps, int total_slots, +__device__ __forceinline__ void tma_copy_gpu_direct_body(GpuDirectStorage const& direct, int warps, int total_slots, unsigned* completed_ctas, uint64_t flag_mc, uint64_t generation) { auto const& c = direct.config; @@ -663,6 +670,83 @@ __global__ void tma_copy_gpu_direct_kernel(GpuDirectStorage direct, int warps, i flag_mc, generation); } +// Scheduler arguments of the fused launch, as halo_q_scheduler_body takes them. +struct FusedSchedulerArgs +{ + int const* routes; + int* out_slots; + int* out_ids; + int* out_levels; + int* out_owners; + unsigned long long const* peer_bases; + int* status; + int* partial; + int* route_aux; + int* grid_sync; + int* plan; + int* route_prefix; + unsigned long long spin_cycles; + int ep; + int experts; + int helpers; + int local_rank; + int route_capacity; + int capacity; + int route_count; +}; + +/* HALO-Q and the GPU-direct copy of its plan in one grid of config.sms CTAs. + * Every CTA first takes its scheduler role (CTA 0 plans, the others write + * routes) and then adds one to route_epoch. Route consumers on other streams + * wait until route_epoch reaches route_target with a stream memory wait, which + * holds no SM. The kernel runs in the remote memory-sync domain while its + * consumers run in the default one, so the release is system-scope. Each CTA + * also waits for route_target before it reads the plan, which CTA 0 finishes + * last, with the scheduler's spin bound. */ +template +__global__ void __launch_bounds__(512, 1) halo_q_tma_fused_kernel(FusedSchedulerArgs const scheduler, + GpuDirectStorage const direct, int warps, int total_slots, unsigned* completed_ctas, uint64_t flag_mc, + uint64_t generation, unsigned long long* route_epoch, unsigned long long route_target) +{ + tensorrt_llm::kernels::halo_q_scheduler_body(scheduler.routes, + scheduler.out_slots, scheduler.out_ids, scheduler.out_levels, scheduler.out_owners, scheduler.peer_bases, + scheduler.status, scheduler.partial, scheduler.route_aux, scheduler.grid_sync, scheduler.plan, + scheduler.route_prefix, nullptr, scheduler.ep, scheduler.experts, scheduler.helpers, scheduler.local_rank, + scheduler.route_capacity, static_cast(gridDim.x), 0, scheduler.spin_cycles, 0, 0, 0, scheduler.capacity, + scheduler.route_count); + // Orders every thread's scheduler writes before thread 0's release. + __syncthreads(); + if (threadIdx.x == 0) + { + unsigned long long observed; + asm volatile("atom.release.sys.global.add.u64 %0, [%1], 1;" : "=l"(observed) : "l"(route_epoch) : "memory"); + long long const start = clock64(); + for (++observed; static_cast(observed - route_target) < 0;) + { + // A CTA that never became resident would hang the copy; fail like a stale plan. + if (clock64() - start > static_cast(scheduler.spin_cycles)) + { + printf("GPU-direct in-switch copy: route epoch timeout; no READY published\n"); + asm volatile("trap;" ::: "memory"); + } + __nanosleep(32); + asm volatile("ld.acquire.sys.global.u64 %0, [%1];" : "=l"(observed) : "l"(route_epoch) : "memory"); + } + } + __syncthreads(); + tma_copy_gpu_direct_body(direct, warps, total_slots, completed_ctas, flag_mc, generation); +} + +// HALO-Q is specialized for 384 experts with 4 helpers at EP4 and EP8. +void const* select_fused_kernel(MegamoeTmaGpuPlanConfig const& config) +{ + if (config.global_experts == 384 && config.helper_count == 4 && config.world == 8) + return reinterpret_cast(&halo_q_tma_fused_kernel<8, 384, 4>); + if (config.global_experts == 384 && config.helper_count == 4 && config.world == 4) + return reinterpret_cast(&halo_q_tma_fused_kernel<4, 384, 4>); + return reinterpret_cast(&halo_q_tma_fused_kernel<0, 0, 0>); +} + } // namespace struct MegamoeTmaCopyState @@ -676,9 +760,82 @@ struct MegamoeTmaCopyState uint64_t* direct_tables = nullptr; bool direct_configured = false; bool direct_bound = false; + // The fused kernel of the bound geometry, prepared by megamoe_tma_copy_bind_gpu_direct. + void const* fused_kernel = nullptr; + int fused_total_slots = 0; + int fused_dynamic_bytes = 0; + cudaLaunchAttribute fused_attributes[2]{}; + unsigned fused_attribute_count = 0; std::mutex submit_mutex; }; +namespace +{ + +// Every fused launch has one HALO-Q CTA of this size per SM. +constexpr int kFusedThreads = 512; + +/* Loads the fused kernel and reserves its local memory at cold setup: doing + * either at the first launch waits for the device, which deadlocks against + * kernels that spin on work enqueued after that launch. Also resolves the slots + * that fit beside both phases' static shared memory; HALO-Q reuses the same + * dynamic allocation before the copy initializes its slots. */ +cudaError_t prepare_fused_launch(MegamoeTmaCopyState& state) +{ + auto const& c = state.config; + auto const& plan = state.direct.config; + void const* kernel = select_fused_kernel(plan); + cudaFuncAttributes attributes{}; + cudaError_t error = cudaFuncGetAttributes(&attributes, kernel); + if (error != cudaSuccess) + return error; + size_t stack = 0; + error = cudaDeviceGetLimit(&stack, cudaLimitStackSize); + if (error == cudaSuccess && stack < attributes.localSizeBytes) + error = cudaDeviceSetLimit(cudaLimitStackSize, attributes.localSizeBytes); + if (error != cudaSuccess) + return error; + int const static_bytes = static_cast(attributes.sharedSizeBytes); + int const slots = megamoe_tma_copy_total_slots(c.device_optin_shared_bytes - static_bytes); + int const capacity = tensorrt_llm::kernels::broadcast_capacity(plan.world, plan.global_experts, plan.helper_count); + int const scheduler_bytes = static_cast(sizeof(int)) + * tensorrt_llm::kernels::halo_q_shared_ints(plan.world, plan.global_experts, capacity, kFusedThreads); + int const slot_bytes = static_cast((slots * (kSlotDataBytes + kSlotControlBytes) + 127) & ~127U); + int const bytes = std::max(slot_bytes, scheduler_bytes); + if (attributes.maxThreadsPerBlock < kFusedThreads || slots < 2 * c.warps + || bytes + static_bytes > c.device_optin_shared_bytes) + return cudaErrorInvalidConfiguration; + error = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes); + if (error == cudaSuccess) + error = cudaFuncSetAttribute( + kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared); + int active = 0; + if (error == cudaSuccess) + error = cudaOccupancyMaxActiveBlocksPerMultiprocessor(&active, kernel, kFusedThreads, bytes); + if (error != cudaSuccess) + return error; + if (active != 1) + return cudaErrorNotSupported; + state.fused_kernel = kernel; + state.fused_total_slots = slots; + state.fused_dynamic_bytes = bytes; + // The copy's NVLink traffic runs in the remote memory-sync domain, so fences + // and completions of kernels in the default domain do not wait for it. + int domains = 0; + if (cudaDeviceGetAttribute(&domains, cudaDevAttrMemSyncDomainCount, c.device) == cudaSuccess && domains >= 2) + { + state.fused_attributes[0].id = cudaLaunchAttributeMemSyncDomain; + state.fused_attributes[0].val.memSyncDomain = cudaLaunchMemSyncDomainRemote; + state.fused_attributes[1].id = cudaLaunchAttributeMemSyncDomainMap; + state.fused_attributes[1].val.memSyncDomainMap.default_ = 0; + state.fused_attributes[1].val.memSyncDomainMap.remote = 1; + state.fused_attribute_count = 2; + } + return cudaSuccess; +} + +} // namespace + extern "C" int megamoe_tma_copy_destroy(MegamoeTmaCopyState** state_ptr) { if (!state_ptr || !*state_ptr) @@ -770,7 +927,7 @@ extern "C" int megamoe_tma_copy_create(uint64_t max_segments, int sms, int warps config.device_sm_count = prop.multiProcessorCount; config.sms = megamoe_tma_copy_sm_count(sms, prop.multiProcessorCount); config.warps = warps ? warps : 7; - config.threads_per_cta = config.warps * MEGAMOE_TMA_COPY_THREADS_PER_WORKER; + config.threads_per_cta = kFusedThreads; config.slice_bytes = kSlice; config.max_segments = max_segments; config.compute_major = prop.major; @@ -785,8 +942,11 @@ extern "C" int megamoe_tma_copy_create(uint64_t max_segments, int sms, int warps &config.device_shared_bytes_per_sm, cudaDevAttrMaxSharedMemoryPerMultiprocessor, config.device); if (error != cudaSuccess) return fail(error); + // The generic fused kernel bounds the geometry; binding prepares the + // specialization of the bound plan geometry. + void const* kernel = reinterpret_cast(&halo_q_tma_fused_kernel<0, 0, 0>); cudaFuncAttributes attributes{}; - error = cudaFuncGetAttributes(&attributes, tma_copy_gpu_direct_kernel); + error = cudaFuncGetAttributes(&attributes, kernel); if (error != cudaSuccess) return fail(error); const size_t available = config.device_optin_shared_bytes > attributes.sharedSizeBytes @@ -807,16 +967,15 @@ extern "C" int megamoe_tma_copy_create(uint64_t max_segments, int sms, int warps config.bank0_slots_per_warp = (config.slots_per_warp + 1) / 2; config.bank1_slots_per_warp = config.slots_per_warp / 2; config.dynamic_shared_bytes = static_cast(dynamic_bytes); - error = cudaFuncSetAttribute( - tma_copy_gpu_direct_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, config.dynamic_shared_bytes); + error = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, config.dynamic_shared_bytes); if (error != cudaSuccess) return fail(error); - error = cudaFuncSetAttribute( - tma_copy_gpu_direct_kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared); + error + = cudaFuncSetAttribute(kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared); if (error != cudaSuccess) return fail(error); - error = cudaOccupancyMaxActiveBlocksPerMultiprocessor(&config.max_active_ctas_per_sm, tma_copy_gpu_direct_kernel, - config.threads_per_cta, config.dynamic_shared_bytes); + error = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &config.max_active_ctas_per_sm, kernel, config.threads_per_cta, config.dynamic_shared_bytes); if (error != cudaSuccess) return fail(error); if (config.max_active_ctas_per_sm != 1) @@ -1010,25 +1169,37 @@ extern "C" int megamoe_tma_copy_bind_gpu_direct(MegamoeTmaCopyState* state, uint d.status = reinterpret_cast(status); d.grid_sync = reinterpret_cast(grid_sync); d.capacity = capacity; + error = prepare_fused_launch(*state); + if (error != cudaSuccess) + return error; state->direct_bound = true; return cudaSuccess; } -extern "C" int megamoe_tma_copy_submit_gpu_direct( - MegamoeTmaCopyState* state, uint64_t flag_mc, uint64_t generation, void* cuda_stream) +int megamoe_tma_copy_submit_fused(MegamoeTmaCopyState* state, + tensorrt_llm::kernels::MoeRebalanceHaloQParams const& scheduler, uint64_t* route_epoch, uint64_t route_target, + uint64_t flag_mc, uint64_t generation, void* cuda_stream) { - if (!state || !flag_mc || (flag_mc & 7) || !generation) + if (!state || !route_epoch || (reinterpret_cast(route_epoch) & 7) || !flag_mc || (flag_mc & 7) + || !generation) return cudaErrorInvalidValue; std::lock_guard lock(state->submit_mutex); if (state->poison != cudaSuccess) return state->poison; - if (!state->direct_bound) + auto const& c = state->config; + auto const& d = state->direct; + // The copy reads the bound scheduler outputs, so the scheduler must write exactly those. + if (!state->direct_bound || scheduler.ctas != c.sms || scheduler.threads != kFusedThreads + || scheduler.planChannel != nullptr || scheduler.outIds != d.ids || scheduler.outLevels != d.levels + || scheduler.outOwners != d.owners || scheduler.planWorkspace != d.workspace || scheduler.status != d.status + || scheduler.gridSync != d.grid_sync || scheduler.ep != d.config.world + || scheduler.experts != d.config.global_experts || scheduler.helpers != d.config.helper_count) return cudaErrorInvalidValue; int device = -1; cudaError_t error = cudaGetDevice(&device); if (error != cudaSuccess) return error; - if (device != state->config.device) + if (device != c.device) return cudaErrorInvalidDevice; auto stream = reinterpret_cast(cuda_stream); cudaStreamCaptureStatus capture; @@ -1037,18 +1208,34 @@ extern "C" int megamoe_tma_copy_submit_gpu_direct( return error; if (capture != cudaStreamCaptureStatusNone) return cudaErrorStreamCaptureUnsupported; - // Raw direct ABI also fixes one stream: protecting descriptor scratch on a - // different stream would not protect HALO-Q's already-enqueued output writes. + // One stream orders the scheduler workspace and the plan scratch. if (state->work_queued && stream != state->last_stream) return cudaErrorInvalidValue; + FusedSchedulerArgs scheduler_args{scheduler.routes, scheduler.outSlots, scheduler.outIds, scheduler.outLevels, + scheduler.outOwners, reinterpret_cast(scheduler.peerBases), scheduler.status, + scheduler.partial, scheduler.routeAux, scheduler.gridSync, scheduler.planWorkspace, scheduler.routePrefix, + static_cast(scheduler.spinCycles), scheduler.ep, scheduler.experts, scheduler.helpers, + scheduler.localRank, scheduler.routeCapacity, + tensorrt_llm::kernels::broadcast_capacity(scheduler.ep, scheduler.experts, scheduler.helpers), + scheduler.validRouteCount}; + GpuDirectStorage direct = d; + int warps = c.warps; + int total_slots = state->fused_total_slots; + unsigned* completed_ctas = state->completed_ctas; + auto* epoch = reinterpret_cast(route_epoch); + unsigned long long target = route_target; + void* arguments[] + = {&scheduler_args, &direct, &warps, &total_slots, &completed_ctas, &flag_mc, &generation, &epoch, &target}; + cudaLaunchConfig_t launch{}; + launch.gridDim = dim3(c.sms, 1, 1); + launch.blockDim = dim3(kFusedThreads, 1, 1); + launch.dynamicSmemBytes = state->fused_dynamic_bytes; + launch.stream = stream; + launch.attrs = state->fused_attributes; + launch.numAttrs = state->fused_attribute_count; state->last_stream = stream; state->work_queued = true; - auto const& c = state->config; - // Keep this an ordinary same-stream launch: programmatic stream serialization - // could observe a stale matching modular epoch before HALO-Q publishes. - tma_copy_gpu_direct_kernel<<>>( - state->direct, c.warps, c.total_slots, state->completed_ctas, flag_mc, generation); - error = cudaGetLastError(); + error = cudaLaunchKernelExC(&launch, state->fused_kernel, arguments); if (error != cudaSuccess) { state->poison = error; diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h index d32029dd7551..bc357698bb7a 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h @@ -183,12 +183,11 @@ extern "C" * Submit is stream-ordered and does not read or wait for plan data on the CPU. * All bound tensors must remain alive and only be reused on the same stream * after the copy kernel has read them. The Python endpoint enforces this lease. + * Binding also loads the fused scheduler + copy kernel. */ int megamoe_tma_copy_configure_gpu_plan(MegamoeTmaCopyState* state, MegamoeTmaGpuPlanConfig const* config); int megamoe_tma_copy_bind_gpu_direct(MegamoeTmaCopyState* state, uint64_t ids, uint64_t levels, uint64_t owners, uint64_t workspace, uint64_t status, uint64_t grid_sync, int capacity); - int megamoe_tma_copy_submit_gpu_direct( - MegamoeTmaCopyState* state, uint64_t flag_mc, uint64_t generation, void* cuda_stream); /* Explicit diagnostic only: synchronizes last work and copies one small result. */ int megamoe_tma_copy_gpu_plan_result(MegamoeTmaCopyState* state, MegamoeTmaGpuPlanResult* result); /* Nulls *state once ownership is consumed. If device activation fails before @@ -199,5 +198,20 @@ extern "C" #ifdef __cplusplus } + +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h" + +/* Launches HALO-Q and the GPU-direct copy of its plan as one kernel of config.sms + * CTAs. scheduler must pass checkMoeRebalanceHaloQParams, use config.sms CTAs and + * no plan channel, and write the outputs bound by megamoe_tma_copy_bind_gpu_direct. + * route_epoch is a caller-owned uint64 that only fused launches change: each CTA + * adds one with a system-scope release once its routes are written, and + * route_target is its value after this launch. A cuStreamWaitValue64 (GEQ + * route_target) on another stream orders that stream after the routes. READY is + * published as by the copy alone. + */ +int megamoe_tma_copy_submit_fused(MegamoeTmaCopyState* state, + tensorrt_llm::kernels::MoeRebalanceHaloQParams const& scheduler, uint64_t* route_epoch, uint64_t route_target, + uint64_t flag_mc, uint64_t generation, void* cuda_stream); #endif #endif diff --git a/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.cpp b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.cpp index c1b82109d806..33b8cfa0855f 100644 --- a/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.cpp +++ b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.cpp @@ -14,7 +14,7 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h" +#include "tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.h" #include #include @@ -51,13 +51,14 @@ int checkedInt(std::int64_t value, char const* name) } // namespace -void moeRebalanceHaloQ(torch::Tensor const& routes, torch::Tensor& outSlots, torch::Tensor& outIds, - torch::Tensor& outLevels, torch::Tensor& outOwners, torch::Tensor const& peerBases, torch::Tensor& status, - torch::Tensor& partial, torch::Tensor& routeAux, torch::Tensor& gridSync, torch::Tensor& planWorkspace, - torch::Tensor& routePrefix, std::int64_t planChannelPtr, std::int64_t epValue, std::int64_t expertsValue, - std::int64_t helpersValue, std::int64_t localRankValue, std::int64_t routeCapacityValue, std::int64_t ctasValue, - std::int64_t threadsValue, bool enablePdl, std::int64_t spinCycles, std::int64_t planAbiVersionValue, - std::int64_t planChannelWordsValue, std::int64_t routeFeaturesValue, std::int64_t validRouteCountValue) +kernels::MoeRebalanceHaloQParams makeMoeRebalanceHaloQParams(torch::Tensor const& routes, torch::Tensor& outSlots, + torch::Tensor& outIds, torch::Tensor& outLevels, torch::Tensor& outOwners, torch::Tensor const& peerBases, + torch::Tensor& status, torch::Tensor& partial, torch::Tensor& routeAux, torch::Tensor& gridSync, + torch::Tensor& planWorkspace, torch::Tensor& routePrefix, std::int64_t planChannelPtr, std::int64_t epValue, + std::int64_t expertsValue, std::int64_t helpersValue, std::int64_t localRankValue, std::int64_t routeCapacityValue, + std::int64_t ctasValue, std::int64_t threadsValue, bool enablePdl, std::int64_t spinCycles, + std::int64_t planAbiVersionValue, std::int64_t planChannelWordsValue, std::int64_t routeFeaturesValue, + std::int64_t validRouteCountValue) { int const ep = checkedInt(epValue, "ep"); int const experts = checkedInt(expertsValue, "experts"); @@ -103,9 +104,7 @@ void moeRebalanceHaloQ(torch::Tensor const& routes, torch::Tensor& outSlots, tor checkCudaTensor(planWorkspace, device, torch::kInt32, 2 + 6 * 384, "plan_workspace"); checkCudaTensor(routePrefix, device, torch::kInt32, static_cast(384) * (ep + 1), "route_prefix"); - c10::cuda::CUDAGuard const deviceGuard(device); - auto const stream = at::cuda::getCurrentCUDAStream(routes.get_device()); - kernels::MoeRebalanceHaloQParams const params{ + return kernels::MoeRebalanceHaloQParams{ routes.data_ptr(), outSlots.data_ptr(), outIds.data_ptr(), @@ -133,7 +132,22 @@ void moeRebalanceHaloQ(torch::Tensor const& routes, torch::Tensor& outSlots, tor routeFeatures, validRouteCount, }; - kernels::invokeMoeRebalanceHaloQ(params, stream); +} + +void moeRebalanceHaloQ(torch::Tensor const& routes, torch::Tensor& outSlots, torch::Tensor& outIds, + torch::Tensor& outLevels, torch::Tensor& outOwners, torch::Tensor const& peerBases, torch::Tensor& status, + torch::Tensor& partial, torch::Tensor& routeAux, torch::Tensor& gridSync, torch::Tensor& planWorkspace, + torch::Tensor& routePrefix, std::int64_t planChannelPtr, std::int64_t ep, std::int64_t experts, + std::int64_t helpers, std::int64_t localRank, std::int64_t routeCapacity, std::int64_t ctas, std::int64_t threads, + bool enablePdl, std::int64_t spinCycles, std::int64_t planAbiVersion, std::int64_t planChannelWords, + std::int64_t routeFeatures, std::int64_t validRouteCount) +{ + auto const params = makeMoeRebalanceHaloQParams(routes, outSlots, outIds, outLevels, outOwners, peerBases, status, + partial, routeAux, gridSync, planWorkspace, routePrefix, planChannelPtr, ep, experts, helpers, localRank, + routeCapacity, ctas, threads, enablePdl, spinCycles, planAbiVersion, planChannelWords, routeFeatures, + validRouteCount); + c10::cuda::CUDAGuard const deviceGuard(routes.device()); + kernels::invokeMoeRebalanceHaloQ(params, at::cuda::getCurrentCUDAStream(routes.get_device())); } } // namespace torch_ext diff --git a/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.h b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.h new file mode 100644 index 000000000000..71213603b6e9 --- /dev/null +++ b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.h @@ -0,0 +1,42 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h" + +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +// Validates the moe_rebalance_halo_q arguments and returns its launch parameters. +// The fused scheduler + in-switch copy op takes the same scheduler arguments. +kernels::MoeRebalanceHaloQParams makeMoeRebalanceHaloQParams(torch::Tensor const& routes, torch::Tensor& outSlots, + torch::Tensor& outIds, torch::Tensor& outLevels, torch::Tensor& outOwners, torch::Tensor const& peerBases, + torch::Tensor& status, torch::Tensor& partial, torch::Tensor& routeAux, torch::Tensor& gridSync, + torch::Tensor& planWorkspace, torch::Tensor& routePrefix, std::int64_t planChannelPtr, std::int64_t ep, + std::int64_t experts, std::int64_t helpers, std::int64_t localRank, std::int64_t routeCapacity, std::int64_t ctas, + std::int64_t threads, bool enablePdl, std::int64_t spinCycles, std::int64_t planAbiVersion, + std::int64_t planChannelWords, std::int64_t routeFeatures, std::int64_t validRouteCount); + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp index d88da5d56031..2e5a9c619f64 100644 --- a/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp +++ b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp @@ -16,6 +16,7 @@ */ #include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h" +#include "tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.h" #include "tensorrt_llm/thop/thUtils.h" #include @@ -228,14 +229,28 @@ class MoeRebalanceTmaState : public torch::CustomClassHolder mBound = true; } - int64_t submitGpuDirect(int64_t flagMc) + // Launches the bound HALO-Q scheduler and this copy as one kernel; see megamoe_tma_copy_submit_fused. + void submitFused(int64_t flagMc, kernels::MoeRebalanceHaloQParams const& scheduler, torch::Tensor const& routeEpoch, + int64_t routeTarget) { std::lock_guard lock(mMutex); requireBound(); + TORCH_CHECK( + routeEpoch.is_cuda() && routeEpoch.get_device() == mDevice, "route_epoch must be on CUDA device ", mDevice); + TORCH_CHECK( + routeEpoch.scalar_type() == at::ScalarType::Long && routeEpoch.is_contiguous() && routeEpoch.numel() == 1, + "route_epoch must be a contiguous one-element torch.int64 tensor"); + uint64_t const flag = checkedUnsigned(flagMc, "flag_mc"); + TORCH_CHECK(flag != 0, "flag_mc must be nonzero"); + kernels::checkMoeRebalanceHaloQParams(scheduler); c10::cuda::CUDAGuard const deviceGuard(c10::Device(c10::DeviceType::CUDA, mDevice)); auto const stream = at::cuda::getCurrentCUDAStream(mDevice).stream(); - submitLocked(flagMc, reinterpret_cast(stream)); - return 0; + checkTmaResult(megamoe_tma_copy_submit_fused(mState, scheduler, + reinterpret_cast(routeEpoch.data_ptr()), + static_cast(routeTarget), flag, mGeneration, reinterpret_cast(stream)), + "moe rebalance fused HALO-Q + TMA submit"); + ++mGeneration; + TORCH_CHECK(mGeneration != 0, "moe rebalance TMA generation overflow"); } int64_t currentGeneration() const @@ -320,16 +335,6 @@ class MoeRebalanceTmaState : public torch::CustomClassHolder return static_cast(value); } - void submitLocked(int64_t flagMc, void* stream) - { - uint64_t const flag = checkedUnsigned(flagMc, "flag_mc"); - TORCH_CHECK(flag != 0, "flag_mc must be nonzero"); - checkTmaResult(megamoe_tma_copy_submit_gpu_direct(mState, flag, mGeneration, stream), - "moe rebalance TMA GPU-direct submit"); - ++mGeneration; - TORCH_CHECK(mGeneration != 0, "moe rebalance TMA generation overflow"); - } - mutable std::mutex mMutex; MegamoeTmaCopyState* mState{nullptr}; int mDevice{-1}; @@ -373,10 +378,20 @@ void moeRebalanceTmaBindGpuDirect(MoeRebalanceTmaStatePtr const& state, torch::T std::move(gridSync), capacity); } -int64_t moeRebalanceTmaSubmitGpuDirect(MoeRebalanceTmaStatePtr const& state, int64_t flagMc) +void moeRebalanceTmaSubmitFused(MoeRebalanceTmaStatePtr const& state, int64_t flagMc, torch::Tensor& routeEpoch, + int64_t routeTarget, torch::Tensor const& routes, torch::Tensor& outSlots, torch::Tensor& outIds, + torch::Tensor& outLevels, torch::Tensor& outOwners, torch::Tensor const& peerBases, torch::Tensor& status, + torch::Tensor& partial, torch::Tensor& routeAux, torch::Tensor& gridSync, torch::Tensor& planWorkspace, + torch::Tensor& routePrefix, int64_t planChannelPtr, int64_t ep, int64_t experts, int64_t helpers, int64_t localRank, + int64_t routeCapacity, int64_t ctas, int64_t threads, bool enablePdl, int64_t spinCycles, int64_t planAbiVersion, + int64_t planChannelWords, int64_t routeFeatures, int64_t validRouteCount) { TORCH_CHECK(state, "moe rebalance TMA state is null"); - return state->submitGpuDirect(flagMc); + auto const scheduler = makeMoeRebalanceHaloQParams(routes, outSlots, outIds, outLevels, outOwners, peerBases, + status, partial, routeAux, gridSync, planWorkspace, routePrefix, planChannelPtr, ep, experts, helpers, + localRank, routeCapacity, ctas, threads, enablePdl, spinCycles, planAbiVersion, planChannelWords, routeFeatures, + validRouteCount); + state->submitFused(flagMc, scheduler, routeEpoch, routeTarget); } int64_t moeRebalanceTmaCurrentGeneration(MoeRebalanceTmaStatePtr const& state) @@ -425,8 +440,14 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state, Tensor ids, Tensor levels, Tensor owners, " "Tensor workspace, Tensor status, Tensor grid_sync, int capacity) -> ()"); m.def( - "moe_rebalance_tma_submit_gpu_direct(" - "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state, int flag_mc) -> int"); + "moe_rebalance_tma_submit_fused(" + "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state, int flag_mc, Tensor(a!) route_epoch, " + "int route_target, Tensor routes, Tensor(b!) out_slots, Tensor(c!) out_ids, Tensor(d!) out_levels, " + "Tensor(e!) out_owners, Tensor peer_bases, Tensor(f!) status, Tensor(g!) partial, Tensor(h!) route_aux, " + "Tensor(i!) grid_sync, Tensor(j!) plan_workspace, Tensor(k!) route_prefix, int plan_channel_ptr, int ep, " + "int experts, int helpers, int local_rank, int route_capacity, int ctas, int threads, bool enable_pdl, " + "int spin_cycles, int plan_abi_version, int plan_channel_words, int route_features, " + "int valid_route_count) -> ()"); m.def( "moe_rebalance_tma_current_gen(" "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state) -> int"); @@ -446,7 +467,7 @@ TORCH_LIBRARY_IMPL(trtllm, CompositeExplicitAutograd, m) m.impl("moe_rebalance_tma_create", &tensorrt_llm::torch_ext::moeRebalanceTmaCreate); m.impl("moe_rebalance_tma_configure_gpu_plan", &tensorrt_llm::torch_ext::moeRebalanceTmaConfigureGpuPlan); m.impl("moe_rebalance_tma_bind_gpu_direct", &tensorrt_llm::torch_ext::moeRebalanceTmaBindGpuDirect); - m.impl("moe_rebalance_tma_submit_gpu_direct", &tensorrt_llm::torch_ext::moeRebalanceTmaSubmitGpuDirect); + m.impl("moe_rebalance_tma_submit_fused", &tensorrt_llm::torch_ext::moeRebalanceTmaSubmitFused); m.impl("moe_rebalance_tma_current_gen", &tensorrt_llm::torch_ext::moeRebalanceTmaCurrentGeneration); m.impl("moe_rebalance_tma_result", &tensorrt_llm::torch_ext::moeRebalanceTmaResult); m.impl("moe_rebalance_tma_config", &tensorrt_llm::torch_ext::moeRebalanceTmaConfig); diff --git a/examples/wide_ep/ep_load_balancer/README.md b/examples/wide_ep/ep_load_balancer/README.md index 96e2987fab29..ca01581c199a 100644 --- a/examples/wide_ep/ep_load_balancer/README.md +++ b/examples/wide_ep/ep_load_balancer/README.md @@ -278,7 +278,7 @@ Offline EPLB keeps one expert placement for the whole run, and online EPLB refre 1. Each rank keeps its own `num_experts / ep_size` experts and has `(num_slots - num_experts) / ep_size` helper slots. The helper slots live in two banks that all MoE layers share, rather than in per-layer replicas. 2. Before a layer's experts run, a GPU scheduler (HALO-Q) exchanges routed-token counts across the EP ranks, picks the hot experts to copy into helper slots, and maps every routed token to a resident or helper slot. -3. A copy kernel moves the picked expert weights over NVLink on a high-priority auxiliary stream, while the main stream runs other work such as input quantization and, for DeepSeek-V4, the shared experts. +3. The same kernel then moves the picked expert weights over NVLink on a high-priority auxiliary stream, while the main stream runs other work such as input quantization and, for DeepSeek-V4, the shared experts. 4. The MoE kernel reads a helper slot only after its weights are marked ready. Expert weights are not staged in host shared memory, so the GDRCopy and shared-memory setup that online EPLB needs does not apply. @@ -308,6 +308,6 @@ Per-layer dynamic EPLB has the following constraints: * **Parallelism**: attention DP (`enable_attention_dp: true`) with pure expert parallelism: `moe_expert_parallel_size` equals the number of GPUs, with no MoE tensor parallelism and no pipeline parallelism. The EP size must be even and between 2 and 32, and the MoE smart router (`moe_cluster_size` greater than 1) is not supported. Unlike `standard` mode, an unsupported model, an EP size of 1, or the smart router raises an error instead of silently disabling the load balancer. * **CUDA graphs**: not supported. A non-null `cuda_graph_config` is replaced with `null` and a warning is logged, and any `prefill_cuda_graph_backend` other than `DISABLED` is rejected. * **Launch queues**: CUDA must be initialized with `CUDA_SCALE_LAUNCH_QUEUES=4x`. The LLM API sets it automatically, in the main process and in every worker, before CUDA is initialized. If the process has already initialized CUDA when the LLM is created, construction fails; export `CUDA_SCALE_LAUNCH_QUEUES=4x` before starting such a process. -* **`auxiliary_sms`**: the number of SMs that the HALO-Q scheduler and the copy kernel use, one after the other. It defaults to 8, must be between 8 and 128, and must be lower than the GPU's SM count. Input quantization and, on SM107, DeepSeek-V4's fused shared-expert kernel leave these SMs free, so planning and copying run alongside them. +* **`auxiliary_sms`**: the number of SMs of the kernel that runs the HALO-Q scheduler and then the copy, one CTA per SM. It defaults to 8, must be between 1 and 8, and must be lower than the GPU's SM count. Input quantization and, on SM107, DeepSeek-V4's fused shared-expert kernel leave these SMs free, so planning and copying run alongside them. * **DWDP**: not supported; `dwdp_config` must keep `dwdp_size` at 1. * **Fields**: only `mode`, `num_slots` and `auxiliary_sms`. `num_slots` is required, must be greater than the expert count, and must be divisible by the EP size. `layer_updates_per_iter` and `initial_global_assignments` are `standard`-only and rejected in this mode. diff --git a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md index 76976420ae35..6679e9ca0f80 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -872,9 +872,11 @@ set wrong has a test that catches it: validation and CUDA initialization/device probes, and propagates it to MPI, direct, and Ray workers. Applications that initialize CUDA before TensorRT-LLM must set it at process startup. -- HALO-Q and the in-switch TMA copy run on the higher-priority copy stream. - The main stream may enqueue shared-expert work and input quantization before - it waits for the plan at the first routed-expert use. A single CPU submitter +- HALO-Q and the in-switch TMA copy run as one kernel on the higher-priority + copy stream, one CTA per auxiliary SM. Each CTA raises a device route epoch + once its routes are written. The main stream may enqueue shared-expert work + and input quantization before its first routed-expert use, where a stream + memory wait on that epoch orders it after the routes. A single CPU submitter owns both streams; the path does not require a CUDA green context or a steady-state host wait. - HALO-Q and TMA follow the native-operator path under diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py index ba46b7025b05..dab7d320b36d 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py @@ -3,14 +3,16 @@ """Single-main-thread HALO-Q and in-switch TMA weight-copy integration. One highest-priority stream is shared across layers. plan_schedule() enqueues -HALO-Q, records route completion, and enqueues TMA copy before returning to the -shared-expert hook. plan_finish() can lend routes before the route event; -the first route consumer calls wait_for_routes(). MegaMoE's -system-acquire READY gate protects helper reads while TMA finishes concurrently. +one kernel that runs HALO-Q and then the TMA copy of its plan on the auxiliary +SMs, before returning to the shared-expert hook. Each CTA of that kernel raises +a device route epoch once its routes are written. plan_finish() can lend routes +before they exist; the first route consumer calls wait_for_routes(), a stream +memory wait on that epoch that holds no SM. MegaMoE's system-acquire READY gate +protects helper reads while the copy finishes concurrently. For each rank and generation g: MAIN: MegaMoE(g) -> consumer_done(g) - COPY: wait consumer_done(g) -> HALO-Q(g+1) EP exchange -> TMA(g+1) + COPY: wait consumer_done(g) -> [HALO-Q(g+1) EP exchange -> TMA(g+1)] All ranks enter the bounded EP exchange after their own previous consumer, so no next-generation overwrite can precede any previous-generation consumer. No host collective or auxiliary copy submitter is needed in the serving path. @@ -27,18 +29,15 @@ __all__ = ["RebalanceSlotSchedulerGroup", "build_rebalance_slot_scheduler_group"] -# Shared-expert persistent grids use the same capacity hint as the copy grid. -TMA_COPY_SM_COUNT = 8 - class _LiveBankLeaseProvider: """Authorize ordered reuse, not a claim of completed GPU work on the host. release_generation_after() enqueues the local consumer wait on COPY. - prove() is legal only after the next HALO-Q launch on that same stream. - Its all-rank count exchange is bounded and precedes plan publication; TMA - waits for that publication and executes after HALO-Q. Therefore a delayed - peer consumer blocks every rank's overwrite without a host-side barrier. + prove() is legal right before the next fused HALO-Q + TMA launch on that + same stream. The launch's bounded all-rank count exchange precedes its plan, + and its copy follows the plan in the same CTAs. Therefore a delayed peer + consumer blocks every rank's overwrite without a host-side barrier. """ collective_safe = True @@ -109,10 +108,6 @@ def __init__( self.topk = int(topk) self.max_tokens_per_rank = int(max_tokens_per_rank) self.auxiliary_sms = int(auxiliary_sms) - if self.auxiliary_sms < TMA_COPY_SM_COUNT: - raise ValueError( - f"auxiliary_sms must be at least {TMA_COPY_SM_COUNT}; got {self.auxiliary_sms}" - ) if copy_stream.device.index != self.device: raise ValueError("copy stream device does not match scheduler device") self.arena = arena @@ -143,8 +138,8 @@ def __init__( logical_expert_count = self.home_experts * ep_size - # Use the EP rank, not the CUDA ordinal. HALO-Q shares the configured - # copy-stream launch budget with TMA. + # Use the EP rank, not the CUDA ordinal. HALO-Q and the TMA copy run in + # the same CTAs, one per auxiliary SM. cfg = CudaSchedulerConfig( ep_size=ep_size, logical_expert_count=logical_expert_count, @@ -164,16 +159,28 @@ def __init__( self._stream_handle = int(self.copy_stream.cuda_stream) self._driver = driver self._input_ready = torch.cuda.Event() - self._plan_ready = torch.cuda.Event() self._consumer_done = torch.cuda.Event() # Materialize each lazily created event once. The steady path uses its # cached driver handle, without per-plan stream guards or allocation. with torch.cuda.device(self.device): - for event in (self._input_ready, self._plan_ready, self._consumer_done): + for event in (self._input_ready, self._consumer_done): event.record(self.copy_stream) self._input_ready_handle = int(self._input_ready.cuda_event) - self._plan_ready_handle = int(self._plan_ready.cuda_event) self._consumer_done_handle = int(self._consumer_done.cuda_event) + error, cu_device = driver.cuDeviceGet(self.device) + supported = 0 + if error == driver.CUresult.CUDA_SUCCESS: + error, supported = driver.cuDeviceGetAttribute( + driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_CAN_USE_64_BIT_STREAM_MEM_OPS, + cu_device, + ) + if error != driver.CUresult.CUDA_SUCCESS or not supported: + raise RuntimeError("MoE rebalance requires 64-bit CUDA stream memory operations") + # CTAs of the fused launches whose routes are written, and its value + # after the latest launch. Only those launches change it. + self._route_epoch = torch.zeros(1, dtype=torch.int64, device=f"cuda:{self.device}") + self._route_epoch_ptr = int(self._route_epoch.data_ptr()) + self._route_target = 0 # Every peer address must be connected before either warmup or submit; # HALO-Q performs an all-rank exchange on the COPY stream. @@ -210,7 +217,7 @@ def __init__( helper_count=self.helper_slots, global_expert_count=logical_expert_count, arena_provider=arena.provider, - tma_sm_count=TMA_COPY_SM_COUNT, + tma_sm_count=self.auxiliary_sms, tma_warps=7, tma_route="plan", ) @@ -320,6 +327,16 @@ def _wait_event(self, stream_handle: int, event_handle: int) -> None: if error != self._driver.CUresult.CUDA_SUCCESS: raise RuntimeError(f"cuStreamWaitEvent(rebalance) failed: {error!r}") + def _wait_route_epoch(self, stream_handle: int) -> None: + (error,) = self._driver.cuStreamWaitValue64( + stream_handle, + self._route_epoch_ptr, + self._route_target, + self._driver.CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_GEQ, + ) + if error != self._driver.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"cuStreamWaitValue64(rebalance) failed: {error!r}") + def plan_schedule(self, logical_expert_ids: torch.Tensor) -> Tuple[Any, int]: """Submit contiguous CUDA int32 routes [tokens, topk] without staging. @@ -348,23 +365,22 @@ def plan_schedule(self, logical_expert_ids: torch.Tensor) -> Tuple[Any, int]: logical_expert_ids.record_stream(self.copy_stream) self._record_event(self._input_ready_handle, self._execution_stream_handle) self._wait_event(self._stream_handle, self._input_ready_handle) - outputs = self.scheduler.submit(logical_expert_ids, self._stream_handle) + outputs = self.scheduler.outputs self._scheduled_generation = self._generation + 1 - self._record_event(self._plan_ready_handle, self._stream_handle) - part = (outputs, num_tokens) - self._plan_part = part if self._pending_release is not None: self.lease.prove(self._pending_release[0]) self._pending_release = None - self._submit_copy(outputs) - return part - - def _submit_copy(self, outputs: object) -> None: - """Consume the HALO-Q plan and enqueue its TMA kernel on COPY.""" - ticket = self.broadcaster.submit(outputs) + route_target = self._route_target + self.auxiliary_sms + ticket = self.broadcaster.submit_fused( + outputs, logical_expert_ids, self._route_epoch, route_target + ) if int(ticket.generation) != self._scheduled_generation: raise RuntimeError("TMA and HALO-Q generation mismatch") + self._route_target = route_target self._generation = int(ticket.generation) + part = (outputs, num_tokens) + self._plan_part = part + return part def plan_finish( self, part: Tuple[Any, int], *, defer_wait: bool = False @@ -387,7 +403,7 @@ def wait_for_routes(self) -> None: if self._generation <= 0: raise RuntimeError("No scheduler generation is available") if self._route_wait_generation != self._generation: - self._wait_event(self._execution_stream_handle, self._plan_ready_handle) + self._wait_route_epoch(self._execution_stream_handle) self._route_wait_generation = self._generation def plan( diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py index cd6ea2d9491b..d9dd5b328e98 100644 --- a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py @@ -410,9 +410,9 @@ def bind_plan_channel( def bind_gpu_direct(self, stream: object) -> None: """Bind a copy consumer to the existing outputs/workspace, without publication. - The consumer must enqueue on this static stream, then call - ``release_plan_channel`` before the next scheduler submission. - No additional plan allocation or GPU/host transfer is performed. + The consumer launches this scheduler with its copy on this static + stream through ``fused_submit_args``. No additional plan allocation or + GPU/host transfer is performed. """ with self._plan_channel_lock: if self._plan_channel_bound: @@ -485,36 +485,44 @@ def _launch_op( state = enter_cuda_stream(stream) try: torch.ops.trtllm.moe_rebalance_halo_q( - routes, - self.outputs.physical_slot_ids, - self.outputs.hot_expert_ids, - self.outputs.hot_expert_group_level, - self.outputs.hot_expert_source_ranks, - self.peer_base, - self.status, - self.partial, - self.route_aux, - self.grid_sync, - self.plan_workspace, - self.route_prefix, - self._plan_channel_ptr, - self.cfg.ep_size, - self.cfg.logical_expert_count, - self.cfg.extra_slots_per_rank, - self.cfg.local_rank, - self.cfg.route_count, - self.cfg.ctas, - self.cfg.threads, - self.cfg.enable_pdl, - SPIN_CYCLES, - self._plan_channel_abi_version, - self._plan_channel_words, - self._plan_channel_route_features, - valid_route_count, + *self._halo_q_args(routes, valid_route_count, enable_pdl=self.cfg.enable_pdl) ) finally: exit_cuda_stream(state) + def _halo_q_args( + self, routes: torch.Tensor, valid_route_count: int, *, enable_pdl: bool + ) -> tuple: + """Positional moe_rebalance_halo_q arguments, shared by the fused copy op.""" + return ( + routes, + self.outputs.physical_slot_ids, + self.outputs.hot_expert_ids, + self.outputs.hot_expert_group_level, + self.outputs.hot_expert_source_ranks, + self.peer_base, + self.status, + self.partial, + self.route_aux, + self.grid_sync, + self.plan_workspace, + self.route_prefix, + self._plan_channel_ptr, + self.cfg.ep_size, + self.cfg.logical_expert_count, + self.cfg.extra_slots_per_rank, + self.cfg.local_rank, + self.cfg.route_count, + self.cfg.ctas, + self.cfg.threads, + enable_pdl, + SPIN_CYCLES, + self._plan_channel_abi_version, + self._plan_channel_words, + self._plan_channel_route_features, + valid_route_count, + ) + def submit( self, routes: torch.Tensor, stream_handle: int, *, valid_tokens: int | None = None, @@ -555,10 +563,40 @@ def submit( self._plan_channel_pending = True return self.outputs + def fused_submit_args(self, routes: torch.Tensor, stream_handle: int) -> tuple: + """Scheduler arguments of the bound GPU-direct copy's fused launch. + + The input and stream contract is ``submit()``'s. The copy endpoint + launches this scheduler together with its copy, so no plan publication + becomes pending and the launch triggers no programmatic dependency. + """ + with self._plan_channel_lock: + if not self._gpu_direct_bound: + raise RuntimeError("a fused launch requires a bound GPU-direct copy") + if stream_handle != self._plan_channel_stream_handle: + raise RuntimeError("bound plan channel requires its static stream") + if self._launch_stream_handle is None: + self._launch_stream_handle = stream_handle + elif stream_handle != self._launch_stream_handle: + raise RuntimeError( + "CUDA scheduler uses one static stream because its workspace " + "and reusable barrier epochs are shared" + ) + valid_tokens = self._check_routes(routes, None) + return self._halo_q_args(routes, valid_tokens * self.cfg.topk, enable_pdl=False) + def _submit_routes( self, routes: torch.Tensor, stream_handle: int, *, valid_tokens: int | None = None, ) -> None: + valid_tokens = self._check_routes(routes, valid_tokens) + self._launch_op( + routes, + stream_handle, + valid_route_count=valid_tokens * self.cfg.topk, + ) + + def _check_routes(self, routes: torch.Tensor, valid_tokens: int | None) -> int: if not isinstance(routes, torch.Tensor): raise TypeError("routes must be a torch.Tensor") if (routes.ndim != 2 or routes.shape[1] != self.cfg.topk @@ -572,11 +610,7 @@ def _submit_routes( valid_tokens = routes.shape[0] if type(valid_tokens) is not int or not 0 <= valid_tokens <= routes.shape[0]: raise ValueError("valid_tokens must be an int in [0,T]") - self._launch_op( - routes, - stream_handle, - valid_route_count=valid_tokens * self.cfg.topk, - ) + return valid_tokens def run( self, diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py index 8669b24d00a2..d8dab8d7d622 100644 --- a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py @@ -478,9 +478,9 @@ def bind_scheduler(self, scheduler: object, stream: object) -> None: if stream_device != scheduler_device: raise ValueError("scheduler stream and output devices differ") bind = getattr(scheduler, "bind_gpu_direct", None) - release = getattr(scheduler, "release_plan_channel", None) - if not callable(bind) or not callable(release): - raise TypeError("scheduler lacks the copy handoff lifecycle") + fused_args = getattr(scheduler, "fused_submit_args", None) + if not callable(bind) or not callable(fused_args): + raise TypeError("scheduler lacks the fused copy launch lifecycle") workspace = getattr(scheduler, "plan_workspace", None) status = getattr(scheduler, "status", None) grid_sync = getattr(scheduler, "grid_sync", None) @@ -616,15 +616,15 @@ def _require_live_bank_reusable(self) -> None: "single live bank is still leased by the previous generation" ) - def submit( - self, - outputs: object, - *, - start_event: object | None = None, - end_event: object | None = None, - payload_end_event: object | None = None, + def submit_fused( + self, outputs: object, routes: torch.Tensor, route_epoch: torch.Tensor, route_target: int ) -> HierarchicalCopyTicket: - """Enqueue the fused GPU-direct TMA copy and READY publication.""" + """Enqueue the bound scheduler and this GPU-direct copy as one kernel. + + Each scheduler CTA adds one to ``route_epoch`` once its routes are + written, which makes it reach ``route_target``; the copy and its READY + publication follow in the same CTAs. + """ with self._lock: if self._handoff_aborted: @@ -635,31 +635,22 @@ def submit( raise RuntimeError("outputs are not from the bound CUDA scheduler") self._require_live_bank_reusable() try: + scheduler_args = self._scheduler.fused_submit_args(routes, self._stream) state = enter_cuda_stream(self._stream_owner) try: - if start_event is not None: - start_event.record() - commands = int( - torch.ops.trtllm.moe_rebalance_tma_submit_gpu_direct( - self.ctx, self.flags_mc - ) + torch.ops.trtllm.moe_rebalance_tma_submit_fused( + self.ctx, self.flags_mc, route_epoch, route_target, *scheduler_args ) - if payload_end_event is not None: - payload_end_event.record() - if end_event is not None: - end_event.record() finally: exit_cuda_stream(state) self._active_generation = self._generation - release = getattr(self._scheduler, "release_plan_channel") - release() except BaseException: self._handoff_aborted = True raise ticket = HierarchicalCopyTicket( generation=self._generation, terminal_flags=self.terminal_flags, - command_count=commands + 1, + command_count=1, ) self._generation += 1 return ticket diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index d490c407fac0..b3dcab68dd76 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1957,10 +1957,10 @@ class MoePerLayerDynamicLoadBalancerConfig(_MoeLoadBalancerConfigBase): status="prototype") auxiliary_sms: int = Field( default=8, - ge=8, - le=128, + ge=1, + le=8, description= - "Number of SMs that the routed-input quantization and shared-expert FC12 kernels leave free for the HALO-Q scheduler and the in-switch copy kernel, which run one after the other on the copy stream. Must be smaller than the device SM count.", + "Number of SMs that the routed-input quantization and shared-expert FC12 kernels leave free for the kernel that runs the HALO-Q scheduler and then the in-switch copy of its plan on the copy stream, one CTA per SM. Must be smaller than the device SM count.", status="prototype") def validate_expert_capacity(self, num_experts: int) -> None: diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index 5d22437c9b32..6c0f63508985 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -4134,9 +4134,9 @@ "properties": { "auxiliary_sms": { "default": 8, - "description": "Number of SMs that the routed-input quantization and shared-expert FC12 kernels leave free for the HALO-Q scheduler and the in-switch copy kernel, which run one after the other on the copy stream. Must be smaller than the device SM count.", - "maximum": 128, - "minimum": 8, + "description": "Number of SMs that the routed-input quantization and shared-expert FC12 kernels leave free for the kernel that runs the HALO-Q scheduler and then the in-switch copy of its plan on the copy stream, one CTA per SM. Must be smaller than the device SM count.", + "maximum": 8, + "minimum": 1, "status": "prototype", "title": "Auxiliary Sms", "type": "integer" diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index 0ce7858017f6..faffa16cef70 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -5698,9 +5698,9 @@ "properties": { "auxiliary_sms": { "default": 8, - "description": "Number of SMs that the routed-input quantization and shared-expert FC12 kernels leave free for the HALO-Q scheduler and the in-switch copy kernel, which run one after the other on the copy stream. Must be smaller than the device SM count.", - "maximum": 128, - "minimum": 8, + "description": "Number of SMs that the routed-input quantization and shared-expert FC12 kernels leave free for the kernel that runs the HALO-Q scheduler and then the in-switch copy of its plan on the copy stream, one CTA per SM. Must be smaller than the device SM count.", + "maximum": 8, + "minimum": 1, "status": "prototype", "title": "Auxiliary Sms", "type": "integer" diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index 84e6e588a65b..1053e5248844 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -173,6 +173,66 @@ class Bundle: assert [pool.bundle for pool in FakePool.instances] == [first_bundle, second_bundle] +def test_hierarchical_fused_submit_launches_bound_scheduler_and_copy(): + source = _TORCH_ROOT / "moe/kernels/megamoe_scheduler/sami/hierarchical.py" + launches = [] + + def submit_fused(ctx, flags_mc, route_epoch, route_target, *scheduler_args): + launches.append((ctx, flags_mc, route_epoch, route_target, scheduler_args)) + + torch = SimpleNamespace( + cuda=SimpleNamespace(device=lambda _: nullcontext(), stream=lambda _: nullcontext()), + ops=SimpleNamespace(trtllm=SimpleNamespace(moe_rebalance_tma_submit_fused=submit_fused)), + ) + ticket_cls = _extract_class( + source, + "HierarchicalCopyTicket", + None, + {"dataclass": dataclasses.dataclass, "field": dataclasses.field}, + ) + broadcast_cls = _extract_class( + source, + "HierarchicalSamiWeightBroadcast", + {"submit_fused", "_require_live_bank_reusable"}, + {"torch": torch, "HierarchicalCopyTicket": ticket_cls}, + ) + outputs = object() + requests = [] + + def fused_submit_args(routes, stream_handle): + requests.append((routes, stream_handle)) + return ("scheduler", routes) + + broadcast = broadcast_cls.__new__(broadcast_cls) + broadcast._lock = nullcontext() + broadcast._handoff_aborted = False + broadcast._outputs = outputs + broadcast._scheduler = SimpleNamespace(fused_submit_args=fused_submit_args) + broadcast._reuse_authority = object() + broadcast._validate_reuse_authority = lambda authority: None + broadcast._active_generation = 0 + broadcast._collective_safe_generation = 0 + broadcast._generation = 1 + broadcast._stream = 19 + broadcast._stream_owner = object() + broadcast.device = 0 + broadcast.ctx = "ctx" + broadcast.flags_mc = 4096 + broadcast.terminal_flags = "terminals" + + ticket = broadcast.submit_fused(outputs, "routes", "epoch", 8) + + assert requests == [("routes", 19)] + assert launches == [("ctx", 4096, "epoch", 8, ("scheduler", "routes"))] + assert (ticket.generation, ticket.command_count) == (1, 1) + assert (broadcast._active_generation, broadcast._generation) == (1, 2) + # The bank stays leased until the consumer of generation 1 is proven done. + with pytest.raises(RuntimeError, match="still leased"): + broadcast.submit_fused(outputs, "routes", "epoch", 16) + with pytest.raises(RuntimeError, match="bound CUDA scheduler"): + broadcast.submit_fused(object(), "routes", "epoch", 16) + + @pytest.mark.parametrize( ( "hidden", @@ -1056,7 +1116,8 @@ class GenerationExecutor: assert args.env_overrides == configured -def test_direct_submit_orders_generations_and_defers_the_route_wait(): +@pytest.mark.parametrize("auxiliary_sms", [4, 8]) +def test_fused_submit_orders_generations_and_defers_the_route_wait(auxiliary_sms): trace = [] state = SimpleNamespace(device=0, thread=7) main = SimpleNamespace(cuda_stream=44, priority=0) @@ -1081,6 +1142,7 @@ def test_direct_submit_orders_generations_and_defers_the_route_wait(): "_finished_generation", "_scheduled_generation", "_route_wait_generation", + "_route_target", "plan_calls", ): setattr(group, name, 0) @@ -1092,12 +1154,15 @@ def test_direct_submit_orders_generations_and_defers_the_route_wait(): group.device = 0 group.topk = 6 group.max_tokens_per_rank = 8192 + group.auxiliary_sms = auxiliary_sms group.copy_stream = group.scheduler_stream = copy group._stream_handle = 19 group._input_ready_handle = 1 - group._plan_ready_handle = 2 group._consumer_done_handle = 3 group._consumer_done = SimpleNamespace(cuda_event=3) + route_epoch = object() + group._route_epoch = route_epoch + group._route_epoch_ptr = 77 def driver_call(*entry): trace.append(entry) @@ -1105,30 +1170,24 @@ def driver_call(*entry): group._driver = SimpleNamespace( CUresult=SimpleNamespace(CUDA_SUCCESS=0), + CUstreamWaitValue_flags=SimpleNamespace(CU_STREAM_WAIT_VALUE_GEQ="geq"), cuEventRecord=lambda event, stream: driver_call("record", event, stream), cuStreamWaitEvent=lambda stream, event, flags: driver_call("wait", stream, event), + cuStreamWaitValue64=lambda stream, address, value, flags: driver_call( + "wait_value", stream, address, value, flags + ), ) outputs = SimpleNamespace(physical_slot_ids=_Routes(8192, trace)) - pending = False - - def submit(routes, handle): - nonlocal pending - assert not pending and handle == 19 - pending = True - trace.append(("HALO", routes.shape[0], id(routes))) - return outputs - - def copy_submit(result): - nonlocal pending - assert pending and result is outputs - pending = False + + def submit_fused(result, routes, epoch, target): + assert result is outputs and epoch is route_epoch generation = group._scheduled_generation - trace.append(("TMA", generation)) + trace.append(("FUSED", routes.shape[0], target, generation)) return SimpleNamespace(generation=generation) - group.scheduler = SimpleNamespace(device="cuda:0", submit=submit) + group.scheduler = SimpleNamespace(device="cuda:0", outputs=outputs) group.broadcaster = SimpleNamespace( - submit=copy_submit, + submit_fused=submit_fused, release_generation_after=lambda generation, event: trace.append( ("consumer_release", generation, event.cuda_event) ), @@ -1139,18 +1198,25 @@ def copy_submit(result): for generation, rows in enumerate((16, 0), 1): trace.clear() routes = _Routes(rows, trace) + target = generation * auxiliary_sms + route_wait = ("wait_value", 44, 77, target, "geq") part = group.plan_schedule(routes) physical, actual_generation = group.plan_finish(part, defer_wait=True) assert physical.shape == (rows, 6) assert actual_generation == generation - assert not any(entry[0] == "wait" and entry[1:] == (44, 2) for entry in trace) + assert route_wait not in trace group.wait_for_routes() group.wait_for_routes() group.finish() - assert sum(entry == ("wait", 44, 2) for entry in trace) == 1 - assert trace.index(("HALO", rows, id(routes))) < trace.index(("TMA", generation)) - assert trace.index(("TMA", generation)) < trace.index(("wait", 44, 2)) + assert trace.count(route_wait) == 1 + assert trace.index(("wait", 19, 1)) < trace.index(("FUSED", rows, target, generation)) + assert trace.index(("FUSED", rows, target, generation)) < trace.index(route_wait) assert trace[-2:] == [("record", 3, 44), ("consumer_release", generation, 3)] + if generation > 1: + # The previous consumer is proven reusable before the next launch. + assert trace.index(("lease", generation - 1)) < trace.index( + ("FUSED", rows, target, generation) + ) # The bound-owner fast path still rejects another thread, stream or device. state.thread = 8 @@ -1174,7 +1240,7 @@ def copy_submit(result): group._driver.cuStreamWaitEvent = lambda *args: (17,) with pytest.raises(RuntimeError, match="cuStreamWaitEvent"): group.plan_schedule(_Routes(8, trace)) - assert not any(entry[0] == "HALO" for entry in trace) + assert not any(entry[0] == "FUSED" for entry in trace) assert group._plan_part is None diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 84046cf7f422..2a38f5b712ec 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -326,21 +326,18 @@ def test_per_layer_dynamic_positive_fields(self, field: str) -> None: with pytest.raises(ValidationError, match=field): MoeConfig(load_balancer=values) - def test_per_layer_dynamic_requires_eight_auxiliary_sms(self) -> None: + @pytest.mark.parametrize("auxiliary_sms", [1, 4]) + def test_per_layer_dynamic_accepts_smaller_auxiliary_sm_budget( + self, auxiliary_sms: int) -> None: values = self._per_layer_dynamic() - values["auxiliary_sms"] = 7 - with pytest.raises(ValidationError, match="auxiliary_sms"): - MoeConfig(load_balancer=values) - - def test_per_layer_dynamic_accepts_larger_auxiliary_sm_budget(self) -> None: - values = self._per_layer_dynamic() - values["auxiliary_sms"] = 12 + values["auxiliary_sms"] = auxiliary_sms config = MoeConfig(load_balancer=values) - assert config.load_balancer.auxiliary_sms == 12 + assert config.load_balancer.auxiliary_sms == auxiliary_sms - def test_per_layer_dynamic_rejects_more_than_halo_q_cta_limit(self) -> None: + def test_per_layer_dynamic_rejects_more_than_fused_copy_sm_limit( + self) -> None: values = self._per_layer_dynamic() - values["auxiliary_sms"] = 129 + values["auxiliary_sms"] = 9 with pytest.raises(ValidationError, match="auxiliary_sms"): MoeConfig(load_balancer=values) diff --git a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py index f25058eddd27..5a76a0c0d105 100644 --- a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py +++ b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py @@ -209,12 +209,12 @@ def test_branch_snapshot_field_types( [ ({"num_slots": 288}, True), ({"mode": "standard", "num_slots": 288, "layer_updates_per_iter": 1}, True), - ({"mode": "per_layer_dynamic", "num_slots": 320, "auxiliary_sms": 16}, True), + ({"mode": "per_layer_dynamic", "num_slots": 320, "auxiliary_sms": 4}, True), # A mapping without mode is standard EPLB, which has no auxiliary_sms. - ({"num_slots": 320, "auxiliary_sms": 16}, False), + ({"num_slots": 320, "auxiliary_sms": 4}, False), ({"mode": "per_layer_dynamic"}, False), ({"mode": "per_layer_dynamic", "num_slots": 320, "layer_updates_per_iter": 1}, False), - ({"mode": "per_layer_dynamic", "num_slots": 320, "auxiliary_sms": 4}, False), + ({"mode": "per_layer_dynamic", "num_slots": 320, "auxiliary_sms": 9}, False), ({"mode": "dynamic", "num_slots": 320}, False), ], ) From a2c457c328562f6b92a3bed3878204a6ba78eddb Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 05:13:18 -0700 Subject: [PATCH 13/59] [None][perf] Keep helper experts resident across forwards in per-layer dynamic EPLB HALO-Q's placement is stable across forwards: a layer's broadcasts often put the same expert into the same helper slot of the same team as in its previous forward, yet every forward copied all of them again over NVLink beside the shared expert. Add moe_config.load_balancer.helper_slot_mode: - per_layer_resident (default): every MoE layer maps its own helper-slot set, and the in-switch copy keeps a per-layer table of the expert each rank's helper slot holds. A broadcast whose whole team already holds its expert in its helper slot moves no data; READY is still published. The last CTA of each copy updates the table from the replicated HALO-Q placement, so every rank keeps the same table. One helper set per layer costs about 6.4 GiB per GPU for DeepSeek-V4-Pro with 3 helper slots per rank. - shared_banks: the previous behavior, two helper banks that all layers alternate between, refilled on every forward. Residency only holds while a layer's own copies are the sole writers of its helper rows, so it is tied to per-layer helper sets. Signed-off-by: jintaop --- .../dynamicEplb/moeRebalancePlaceholder.cpp | 2 +- .../dynamicEplb/moeRebalanceTma.cu | 76 ++++++++++++++++--- .../loadBalance/dynamicEplb/moeRebalanceTma.h | 7 +- .../moeRebalanceTmaGpuPlanBuilder.cuh | 25 +++++- .../moe/loadBalance/moeRebalanceTmaOp.cpp | 11 +-- examples/wide_ep/README.md | 4 +- examples/wide_ep/ep_load_balancer/README.md | 8 +- .../moe/fused_moe/MOE_DEVELOPER_GUIDE.md | 4 + .../mega_moe/rebalance_live_arena.py | 52 ++++++++----- .../mega_moe/rebalance_slot_scheduler.py | 8 +- .../moe/fused_moe/per_layer_dynamic_eplb.py | 7 +- .../megamoe_scheduler/sami/hierarchical.py | 14 +++- tensorrt_llm/llmapi/llm_args.py | 7 +- .../schemas/trtllm-serve-config.schema.json | 13 +++- .../trtllm-serve-disagg-config.schema.json | 13 +++- .../usage/llm_args_golden_manifest.json | 9 +++ .../_torch/moe/test_dynamic_eplb_contracts.py | 48 ++++++++++++ .../_torch/moe/test_per_layer_dynamic_eplb.py | 22 ++++++ tests/unittest/llmapi/test_llm_args.py | 24 ++++++ .../test_trtllm_serve_config_schemas.py | 5 ++ 20 files changed, 307 insertions(+), 52 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlaceholder.cpp b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlaceholder.cpp index 40ca0ed5ae89..259820957bf9 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlaceholder.cpp +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlaceholder.cpp @@ -69,7 +69,7 @@ extern "C" int megamoe_tma_copy_configure_gpu_plan(MegamoeTmaCopyState*, Megamoe } extern "C" int megamoe_tma_copy_bind_gpu_direct( - MegamoeTmaCopyState*, uint64_t, uint64_t, uint64_t, uint64_t, uint64_t, uint64_t, int) + MegamoeTmaCopyState*, uint64_t, uint64_t, uint64_t, uint64_t, uint64_t, uint64_t, int, int) { return unsupported(); } diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu index 1d5d4dfdd418..a2047a3b9ca5 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu @@ -193,7 +193,40 @@ __device__ __forceinline__ void multicast_generation(uint64_t flag_mc, uint64_t asm volatile("fence.proxy.alias;" ::: "memory"); } -__device__ __forceinline__ void finish_cta_and_notify(unsigned* completed_ctas, uint64_t flag_mc, uint64_t generation) +/* Helper residency of one layer: the expert that each rank's helper slot holds, + * int32 [world][helper_count], -1 when unknown. Every rank keeps an identical + * table because it changes only from the replicated HALO-Q placement. */ +struct HelperResidency +{ + int* table; + int const* placement; // HALO-Q workspace: epoch, count, six capacity-strided fields. + int capacity; + int helpers; + int world; +}; + +// After a copy, every member of each broadcast's team holds its expert in the +// broadcast's helper slot. +__device__ __forceinline__ void update_helper_residency(HelperResidency const& residency) +{ + int const count = __ldcg(residency.placement + 1); + int const used = count < 0 ? 0 : (count > residency.capacity ? residency.capacity : count); + int const* fields = residency.placement + 2; + for (int broadcast = 0; broadcast < used; ++broadcast) + { + int const expert = __ldcg(fields + broadcast); + int const begin = __ldcg(fields + 2 * residency.capacity + broadcast); + int const size = __ldcg(fields + 3 * residency.capacity + broadcast); + int const helper = __ldcg(fields + 5 * residency.capacity + broadcast); + if (helper < 0 || helper >= residency.helpers || begin < 0 || size <= 0 || begin + size > residency.world) + continue; + for (int rank = begin; rank < begin + size; ++rank) + residency.table[rank * residency.helpers + helper] = expert; + } +} + +__device__ __forceinline__ void finish_cta_and_notify( + unsigned* completed_ctas, uint64_t flag_mc, uint64_t generation, HelperResidency const* residency = nullptr) { unsigned previous; // A release/acquire RMW chain combines the completed payload of all CTAs. @@ -205,6 +238,10 @@ __device__ __forceinline__ void finish_cta_and_notify(unsigned* completed_ctas, // the last CTA can rearm the counter without a host memset each time. asm volatile("st.relaxed.gpu.global.u32 [%0], 0;" ::"l"(completed_ctas) : "memory"); multicast_generation(flag_mc, generation); + // Every CTA read the table before joining the chain; the next copy of + // this layer is a later kernel on the same stream. + if (residency != nullptr) + update_helper_residency(*residency); } } @@ -483,7 +520,7 @@ __device__ __forceinline__ void publish_bank_ticket(unsigned* address, unsigned __device__ __forceinline__ void tma_copy_body(MegamoeTmaCopySegment const* segments, MegamoeTmaCopyRange const* ranges, uint64_t range_count, uint64_t total_slices, int warps, int total_slots, unsigned* completed_ctas, uint64_t flag_mc, - uint64_t generation) + uint64_t generation, HelperResidency const* residency) { if (!total_slices) { @@ -592,7 +629,7 @@ __device__ __forceinline__ void tma_copy_body(MegamoeTmaCopySegment const* segme } __syncthreads(); if (threadIdx.x == 0 && flag_mc) - finish_cta_and_notify(completed_ctas, flag_mc, generation); + finish_cta_and_notify(completed_ctas, flag_mc, generation, residency); } // Per-CTA scratch is allocated once. No plan tensor is copied between kernels. @@ -606,6 +643,8 @@ struct GpuDirectStorage MegamoeTmaCopySegment* segments; MegamoeTmaCopyRange* ranges; MegamoeTmaGpuPlanResult* results; + // Helper residency table, or null when helper slots are shared across layers. + int* resident; }; __device__ __forceinline__ void tma_copy_gpu_direct_body(GpuDirectStorage const& direct, int warps, int total_slots, @@ -647,8 +686,8 @@ __device__ __forceinline__ void tma_copy_gpu_direct_body(GpuDirectStorage const& __syncthreads(); if (threadIdx.x == 0) { - megamoe_tma_build_gpu_plan( - &c, plan, plan + c.plan_words, segments, direct.max_segments, ranges, direct.max_segments, result); + megamoe_tma_build_gpu_plan(&c, plan, plan + c.plan_words, segments, direct.max_segments, ranges, + direct.max_segments, result, direct.resident); // Malformed plans must never publish READY. The CUDA error surfaces at // the caller's normal completion check, without a host wait in submit. if (result->error) @@ -658,16 +697,18 @@ __device__ __forceinline__ void tma_copy_gpu_direct_body(GpuDirectStorage const& } } __syncthreads(); + HelperResidency const residency{direct.resident, direct.workspace, cap, c.helper_count, c.world}; + HelperResidency const* const update = direct.resident != nullptr ? &residency : nullptr; if (!result->total_slices) { // Work size is device-derived, so the launch always has config.sms CTAs. // Empty ranks must still join the last-CTA chain before one READY store. if (threadIdx.x == 0) - finish_cta_and_notify(completed_ctas, flag_mc, generation); + finish_cta_and_notify(completed_ctas, flag_mc, generation, update); return; } tma_copy_body(segments, ranges, result->range_count, result->total_slices, warps, total_slots, completed_ctas, - flag_mc, generation); + flag_mc, generation, update); } // Scheduler arguments of the fused launch, as halo_q_scheduler_body takes them. @@ -871,6 +912,8 @@ extern "C" int megamoe_tma_copy_destroy(MegamoeTmaCopyState** state_ptr) record(cudaFree(state->direct.ranges)); if (state->direct.results) record(cudaFree(state->direct.results)); + if (state->direct.resident) + record(cudaFree(state->direct.resident)); if (state->completed_ctas) record(cudaFree(state->completed_ctas)); if (previous_device != state->config.device) @@ -1138,10 +1181,11 @@ extern "C" int megamoe_tma_copy_configure_gpu_plan(MegamoeTmaCopyState* state, M } extern "C" int megamoe_tma_copy_bind_gpu_direct(MegamoeTmaCopyState* state, uint64_t ids, uint64_t levels, - uint64_t owners, uint64_t workspace, uint64_t status, uint64_t grid_sync, int capacity) + uint64_t owners, uint64_t workspace, uint64_t status, uint64_t grid_sync, int capacity, int helper_residency) { if (!state || !ids || !levels || !owners || !workspace || !status || !grid_sync - || (ids | levels | owners | workspace | status | grid_sync) % alignof(int) || capacity < 1 || capacity > 384) + || (ids | levels | owners | workspace | status | grid_sync) % alignof(int) || capacity < 1 || capacity > 384 + || (helper_residency != 0 && helper_residency != 1)) return cudaErrorInvalidValue; std::lock_guard lock(state->submit_mutex); if (!state->direct_configured || state->direct_bound || state->work_queued) @@ -1169,6 +1213,18 @@ extern "C" int megamoe_tma_copy_bind_gpu_direct(MegamoeTmaCopyState* state, uint d.status = reinterpret_cast(status); d.grid_sync = reinterpret_cast(grid_sync); d.capacity = capacity; + if (helper_residency) + { + size_t const bytes + = sizeof(int) * static_cast(d.config.world) * static_cast(d.config.helper_count); + error = cudaMalloc(reinterpret_cast(&d.resident), bytes); + if (error == cudaSuccess) + error = cudaMemset(d.resident, 0xFF, bytes); // -1: no expert is known to be resident. + if (error == cudaSuccess) + error = cudaStreamSynchronize(nullptr); // Cold initialization only. + if (error != cudaSuccess) + return error; + } error = prepare_fused_launch(*state); if (error != cudaSuccess) return error; @@ -1208,7 +1264,7 @@ int megamoe_tma_copy_submit_fused(MegamoeTmaCopyState* state, return error; if (capture != cudaStreamCaptureStatusNone) return cudaErrorStreamCaptureUnsupported; - // One stream orders the scheduler workspace and the plan scratch. + // One stream orders the scheduler workspace, the plan scratch and the residency table. if (state->work_queued && stream != state->last_stream) return cudaErrorInvalidValue; FusedSchedulerArgs scheduler_args{scheduler.routes, scheduler.outSlots, scheduler.outIds, scheduler.outLevels, diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h index bc357698bb7a..dffa0e523298 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h @@ -183,11 +183,14 @@ extern "C" * Submit is stream-ordered and does not read or wait for plan data on the CPU. * All bound tensors must remain alive and only be reused on the same stream * after the copy kernel has read them. The Python endpoint enforces this lease. - * Binding also loads the fused scheduler + copy kernel. + * Binding also loads the fused scheduler + copy kernel. helper_residency=1 + * keeps a per-layer table of the expert each helper slot holds and skips a + * broadcast whose whole team already holds its expert; it requires helper + * rows that only this copy writes. */ int megamoe_tma_copy_configure_gpu_plan(MegamoeTmaCopyState* state, MegamoeTmaGpuPlanConfig const* config); int megamoe_tma_copy_bind_gpu_direct(MegamoeTmaCopyState* state, uint64_t ids, uint64_t levels, uint64_t owners, - uint64_t workspace, uint64_t status, uint64_t grid_sync, int capacity); + uint64_t workspace, uint64_t status, uint64_t grid_sync, int capacity, int helper_residency); /* Explicit diagnostic only: synchronizes last work and copies one small result. */ int megamoe_tma_copy_gpu_plan_result(MegamoeTmaCopyState* state, MegamoeTmaGpuPlanResult* result); /* Nulls *state once ownership is consumed. If device activation fails before diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlanBuilder.cuh b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlanBuilder.cuh index 351bc18774b7..6441e46a10fe 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlanBuilder.cuh +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlanBuilder.cuh @@ -33,9 +33,13 @@ #endif namespace megamoe_gpu_plan_detail { +/* A helper slot whose whole team already holds its expert needs no copy. */ +static constexpr int kSlotModeResident = 3; + struct Builder : MegamoeTmaGpuPlanConfig { int32_t const* plan; + int32_t const* resident; /* [world][helper_count] expert held per helper slot, or null. */ int32_t* slot_scratch; MegamoeTmaCopySegment* tma_segments; MegamoeTmaCopyRange* tma_ranges; @@ -144,6 +148,16 @@ static MEGAMOE_GP_HD inline int append_tma_range(Builder* c, size_t* ranges, uin return 0; } +static MEGAMOE_GP_HD inline bool team_holds_expert(Builder const* c, int slot, int expert, int level) +{ + int group = c->group_sizes[level]; + int begin = c->rank / group * group; + for (int rank = begin; rank < begin + group; ++rank) + if (c->resident[rank * c->helper_count + slot] != expert) + return false; + return true; +} + static MEGAMOE_GP_HD inline size_t build_tma_descriptors(Builder* c, int outgoing_count, int* plan_error) { *plan_error = 3; @@ -161,6 +175,11 @@ static MEGAMOE_GP_HD inline size_t build_tma_descriptors(Builder* c, int outgoin { if (experts[slot] < 0 || modes[slot] == 2) continue; + if (c->resident && team_holds_expert(c, slot, experts[slot], levels[slot])) + { + modes[slot] = kSlotModeResident; + continue; + } if (c->tma_route == 1) { int group = c->group_sizes[levels[slot]]; @@ -424,10 +443,13 @@ static MEGAMOE_GP_HD inline size_t build_descriptors(Builder* c, int* plan_error /* Host/device parity permits differential tests against the CPU submitter. * Counts remain zero on failure: callers must not publish successful READY. * Caller supplies valid memory for the full immutable tables and scratch. + * A non-null resident table removes the slots whose team already holds the + * planned expert from every routing domain. */ static MEGAMOE_GP_HD inline int megamoe_tma_build_gpu_plan(MegamoeTmaGpuPlanConfig const* config, int32_t const* device_plan, int32_t* slot_scratch, MegamoeTmaCopySegment* segments, uint64_t max_segments, - MegamoeTmaCopyRange* ranges, uint64_t max_ranges, MegamoeTmaGpuPlanResult* result) + MegamoeTmaCopyRange* ranges, uint64_t max_ranges, MegamoeTmaGpuPlanResult* result, + int32_t const* resident = nullptr) { if (!result) return 1; @@ -463,6 +485,7 @@ static MEGAMOE_GP_HD inline int megamoe_tma_build_gpu_plan(MegamoeTmaGpuPlanConf megamoe_gpu_plan_detail::Builder c{}; static_cast(c) = g; c.plan = device_plan; + c.resident = resident; c.slot_scratch = slot_scratch; c.tma_segments = segments; c.tma_ranges = ranges; diff --git a/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp index 2e5a9c619f64..c81984897757 100644 --- a/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp +++ b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp @@ -195,7 +195,7 @@ class MoeRebalanceTmaState : public torch::CustomClassHolder } void bindGpuDirect(torch::Tensor ids, torch::Tensor levels, torch::Tensor owners, torch::Tensor workspace, - torch::Tensor status, torch::Tensor gridSync, int64_t capacity) + torch::Tensor status, torch::Tensor gridSync, int64_t capacity, bool helperResidency) { std::lock_guard lock(mMutex); requireLive(); @@ -215,7 +215,7 @@ class MoeRebalanceTmaState : public torch::CustomClassHolder megamoe_tma_copy_bind_gpu_direct(mState, reinterpret_cast(ids.data_ptr()), reinterpret_cast(levels.data_ptr()), reinterpret_cast(owners.data_ptr()), reinterpret_cast(workspace.data_ptr()), reinterpret_cast(status.data_ptr()), - reinterpret_cast(gridSync.data_ptr()), capacityValue), + reinterpret_cast(gridSync.data_ptr()), capacityValue, helperResidency ? 1 : 0), "moe rebalance TMA GPU-direct bind"); // The native state stores raw device addresses. Retaining the tensors in @@ -371,11 +371,12 @@ void moeRebalanceTmaConfigureGpuPlan(MoeRebalanceTmaStatePtr const& state, int64 } void moeRebalanceTmaBindGpuDirect(MoeRebalanceTmaStatePtr const& state, torch::Tensor ids, torch::Tensor levels, - torch::Tensor owners, torch::Tensor workspace, torch::Tensor status, torch::Tensor gridSync, int64_t capacity) + torch::Tensor owners, torch::Tensor workspace, torch::Tensor status, torch::Tensor gridSync, int64_t capacity, + bool helperResidency) { TORCH_CHECK(state, "moe rebalance TMA state is null"); state->bindGpuDirect(std::move(ids), std::move(levels), std::move(owners), std::move(workspace), std::move(status), - std::move(gridSync), capacity); + std::move(gridSync), capacity, helperResidency); } void moeRebalanceTmaSubmitFused(MoeRebalanceTmaStatePtr const& state, int64_t flagMc, torch::Tensor& routeEpoch, @@ -438,7 +439,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) m.def( "moe_rebalance_tma_bind_gpu_direct(" "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state, Tensor ids, Tensor levels, Tensor owners, " - "Tensor workspace, Tensor status, Tensor grid_sync, int capacity) -> ()"); + "Tensor workspace, Tensor status, Tensor grid_sync, int capacity, bool helper_residency) -> ()"); m.def( "moe_rebalance_tma_submit_fused(" "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state, int flag_mc, Tensor(a!) route_epoch, " diff --git a/examples/wide_ep/README.md b/examples/wide_ep/README.md index 5770e8a29681..2973ce4db0ad 100644 --- a/examples/wide_ep/README.md +++ b/examples/wide_ep/README.md @@ -82,7 +82,7 @@ Attention DP (`enable_attention_dp: true`) is not part of this gate, so EPLB is Load-balancing mode, `standard` or `per_layer_dynamic`. Configs without `mode` are treated as `standard`. * `standard`: offline or online EPLB, selected by [`layer_updates_per_iter`](#layer_updates_per_iter). -* `per_layer_dynamic`: per-layer dynamic EPLB, which re-plans the expert replicas of every MoE layer on every forward pass from that pass's routing. It requires the `MEGAMOE_CUTEDSL` backend (`AUTO` resolves to it) and accepts only `num_slots` and `auxiliary_sms` besides `mode`. See [Per-Layer Dynamic EP Load Balancer](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/wide_ep/ep_load_balancer#per-layer-dynamic-ep-load-balancer) for an example and its constraints. +* `per_layer_dynamic`: per-layer dynamic EPLB, which re-plans the expert replicas of every MoE layer on every forward pass from that pass's routing. It requires the `MEGAMOE_CUTEDSL` backend (`AUTO` resolves to it) and accepts only `num_slots`, `auxiliary_sms` and `helper_slot_mode` besides `mode`. See [Per-Layer Dynamic EP Load Balancer](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/wide_ep/ep_load_balancer#per-layer-dynamic-ep-load-balancer) for an example and its constraints. #### `num_slots` @@ -92,7 +92,7 @@ Total number of expert slots. It must be ≥ total experts, and divisible by the 2. Set to number of total experts + EP size, such as 288 for DeepSeek R1, 32-way EP. * This means there is 1 extra expert on each EP rank, so that there is more room for the per-rank token distribution to be more balanced. -With `mode: per_layer_dynamic`, `num_slots` is required and must be greater than the total expert count, which must itself be divisible by the EP size. Each rank keeps its `num_experts / ep_size` experts and gets `(num_slots - num_experts) / ep_size` extra slots. Unlike `standard` mode, the extra slots are shared helper banks rather than per-layer replicas: all MoE layers alternate between the same two helper banks, and each layer refills its bank with copies of its hot experts on every forward pass. For example, 384 experts with `num_slots: 416` on 8 EP ranks give each rank 48 experts and 4 helper slots. +With `mode: per_layer_dynamic`, `num_slots` is required and must be greater than the total expert count, which must itself be divisible by the EP size. Each rank keeps its `num_experts / ep_size` experts and gets `(num_slots - num_experts) / ep_size` extra slots. Unlike `standard` mode, the extra slots are helper slots rather than static replicas: each layer refills them with copies of its hot experts on every forward pass. By default (`helper_slot_mode: per_layer_resident`) each layer has its own helper slots and skips experts they already hold; `shared_banks` has all MoE layers alternate between two shared helper banks. For example, 384 experts with `num_slots: 416` on 8 EP ranks give each rank 48 experts and 4 helper slots. To disable MoE load balancing, omit the whole `load_balancer` section. EPLB is enabled by the presence of `load_balancer`, so setting `num_slots` to 0 does not disable it. diff --git a/examples/wide_ep/ep_load_balancer/README.md b/examples/wide_ep/ep_load_balancer/README.md index ca01581c199a..850115bf8b0a 100644 --- a/examples/wide_ep/ep_load_balancer/README.md +++ b/examples/wide_ep/ep_load_balancer/README.md @@ -278,7 +278,7 @@ Offline EPLB keeps one expert placement for the whole run, and online EPLB refre 1. Each rank keeps its own `num_experts / ep_size` experts and has `(num_slots - num_experts) / ep_size` helper slots. The helper slots live in two banks that all MoE layers share, rather than in per-layer replicas. 2. Before a layer's experts run, a GPU scheduler (HALO-Q) exchanges routed-token counts across the EP ranks, picks the hot experts to copy into helper slots, and maps every routed token to a resident or helper slot. -3. The same kernel then moves the picked expert weights over NVLink on a high-priority auxiliary stream, while the main stream runs other work such as input quantization and, for DeepSeek-V4, the shared experts. +3. The same kernel then moves the picked expert weights over NVLink on a high-priority auxiliary stream, while the main stream runs other work such as input quantization and, for DeepSeek-V4, the shared experts. By default each layer keeps its own helper slots, and an expert that a helper slot already holds from the previous forward pass is not copied again. 4. The MoE kernel reads a helper slot only after its weights are marked ready. Expert weights are not staged in host shared memory, so the GDRCopy and shared-memory setup that online EPLB needs does not apply. @@ -295,6 +295,7 @@ moe_config: mode: per_layer_dynamic num_slots: 416 # 384 experts + 4 helper slots on each of 8 EP ranks auxiliary_sms: 8 # optional; 8 is the default + helper_slot_mode: per_layer_resident # optional; the default EOF ``` @@ -304,10 +305,11 @@ Per-layer dynamic EPLB has the following constraints: * **Backend**: `MEGAMOE_CUTEDSL` only. `AUTO`, the default, resolves to it; any other explicit backend is rejected. * **Models**: an architecture listed in `moe_model_arch_list` (see [`load_balancer`](../README.md#load_balancer)) with NVFP4 MoE weights, as `MEGAMOE_CUTEDSL` requires. The model can have at most 384 routed experts, and the expert count must be divisible by the EP size. -* **GPUs**: SM100, SM103 or SM107, with all EP ranks in one NVLink domain that supports multicast. The helper banks are CUDA fabric-handle allocations bound to NVLink multicast objects. +* **GPUs**: SM100, SM103 or SM107, with all EP ranks in one NVLink domain that supports multicast. The helper slots are CUDA fabric-handle allocations bound to NVLink multicast objects. * **Parallelism**: attention DP (`enable_attention_dp: true`) with pure expert parallelism: `moe_expert_parallel_size` equals the number of GPUs, with no MoE tensor parallelism and no pipeline parallelism. The EP size must be even and between 2 and 32, and the MoE smart router (`moe_cluster_size` greater than 1) is not supported. Unlike `standard` mode, an unsupported model, an EP size of 1, or the smart router raises an error instead of silently disabling the load balancer. * **CUDA graphs**: not supported. A non-null `cuda_graph_config` is replaced with `null` and a warning is logged, and any `prefill_cuda_graph_backend` other than `DISABLED` is rejected. * **Launch queues**: CUDA must be initialized with `CUDA_SCALE_LAUNCH_QUEUES=4x`. The LLM API sets it automatically, in the main process and in every worker, before CUDA is initialized. If the process has already initialized CUDA when the LLM is created, construction fails; export `CUDA_SCALE_LAUNCH_QUEUES=4x` before starting such a process. * **`auxiliary_sms`**: the number of SMs of the kernel that runs the HALO-Q scheduler and then the copy, one CTA per SM. It defaults to 8, must be between 1 and 8, and must be lower than the GPU's SM count. Input quantization and, on SM107, DeepSeek-V4's fused shared-expert kernel leave these SMs free, so planning and copying run alongside them. +* **`helper_slot_mode`**: `per_layer_resident` (the default) gives every MoE layer its own helper slots and skips copying experts that they still hold from the previous forward pass. It costs one helper set per MoE layer, about 6.4 GiB per GPU for DeepSeek-V4-Pro with 3 helper slots per rank. `shared_banks` has all MoE layers alternate between two shared helper banks and copies every helper expert in every forward pass. * **DWDP**: not supported; `dwdp_config` must keep `dwdp_size` at 1. -* **Fields**: only `mode`, `num_slots` and `auxiliary_sms`. `num_slots` is required, must be greater than the expert count, and must be divisible by the EP size. `layer_updates_per_iter` and `initial_global_assignments` are `standard`-only and rejected in this mode. +* **Fields**: only `mode`, `num_slots`, `auxiliary_sms` and `helper_slot_mode`. `num_slots` is required, must be greater than the expert count, and must be divisible by the EP size. `layer_updates_per_iter` and `initial_global_assignments` are `standard`-only and rejected in this mode. diff --git a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md index 6679e9ca0f80..5e3eb29e1185 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -879,6 +879,10 @@ set wrong has a test that catches it: memory wait on that epoch orders it after the routes. A single CPU submitter owns both streams; the path does not require a CUDA green context or a steady-state host wait. +- `helper_slot_mode=per_layer_resident` gives each layer its own helper-slot + set, and the copy keeps a per-layer table of the expert each helper slot + holds so that it skips experts a whole team already holds. Any other writer + of those helper rows would invalidate the table. - HALO-Q and TMA follow the native-operator path under `cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb` and `cpp/tensorrt_llm/thop/moe/loadBalance`. The scheduler accepts the original diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_live_arena.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_live_arena.py index bb1f0a025873..81e8655615c8 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_live_arena.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_live_arena.py @@ -2,11 +2,12 @@ # SPDX-License-Identifier: Apache-2.0 """EP-scoped live-weight storage for HALO-Q and hierarchical SAMI copies. -Every layer owns its home-expert storage and maps one of two shared helper-slot -banks directly after it in virtual address space. Layers alternate banks, while -per-layer READY terminals and small planes remain independent. Allocation and -binding are collective across the EP group; vendor imports stay lazy until the -rebalance path is enabled. +Every layer owns its home-expert storage and maps a helper-slot set directly +after it in virtual address space: either its own set, which lets helper experts +stay resident across forwards, or one of two banks that all layers alternate +between. Per-layer READY terminals and small planes remain independent. +Allocation and binding are collective across the EP group; vendor imports stay +lazy until the rebalance path is enabled. """ import math @@ -525,13 +526,15 @@ def __init__(self, pool: "_SharedSlotPool", helper_set: int) -> None: class _SharedSlotPool: - """Helper-slot sets shared by every MoE layer of one geometry. + """Helper-slot sets of every MoE layer of one geometry. Layer ``i`` maps helper set ``i % sets`` directly after its home rows, so - each plane is still one contiguous ``[H + S]`` kernel view. Only the home - rows are per layer. No extra synchronization guards reuse: layer ``i``'s - copy follows HALO-Q's all-rank exchange, which every rank enters only after - MAIN finished the previous layers' MegaMoE. + each plane is still one contiguous ``[H + S]`` kernel view. With + ``per_layer_sets`` the pool adds one set per layer as the layers are built, + so only a layer's own copies write its helper rows; otherwise every layer + shares two sets. No extra synchronization guards reuse of a shared set: + layer ``i``'s copy follows HALO-Q's all-rank exchange, which every rank + enters only after MAIN finished the previous layers' MegaMoE. READY terminals stay per layer because one forward shares its generation across layers. They and the fp32 planes live in per-layer records. @@ -548,6 +551,7 @@ def __init__( device: int, comm: Any, owner_registry: List["_SharedSlotPool"], + per_layer_sets: bool, ) -> None: from tensorrt_llm._torch.moe.kernels.megamoe_scheduler.sami import hierarchy_group_sizes @@ -600,10 +604,12 @@ def __init__( self.pages: List[_RecordPage] = [] self.layer_count = 0 self.sets: Tuple[_HelperSet, ...] = () + self.per_layer_sets = per_layer_sets self._construction_complete = False self._owner_registry = owner_registry self._owner_registry.append(self) - self.sets = self._create_sets(comm) + if not per_layer_sets: + self.sets = self._create_sets(_HELPER_BANK_COUNT, comm) self._construction_complete = True def register_producer(self, producer: Any) -> None: @@ -715,14 +721,11 @@ def close_collectively(self, local_safe: bool = True) -> bool: finally: self._closing = False - def _create_sets(self, comm: Any) -> Tuple[_HelperSet, ...]: + def _create_sets(self, count: int, comm: Any) -> Tuple[_HelperSet, ...]: if not self.large: - return tuple( - _HelperSet({}, tuple({} for _ in self.groups), ()) - for _ in range(_HELPER_BANK_COUNT) - ) + return tuple(_HelperSet({}, tuple({} for _ in self.groups), ()) for _ in range(count)) sets, specs = [], [] - for _ in range(_HELPER_BANK_COUNT): + for _ in range(count): chunks = {p.name: self.vmm.create(self.helper_bytes[p.name])[0] for p in self.large} level_bases, level_specs = [], [] for members in self.groups: @@ -801,6 +804,8 @@ def build_layer(self, comm: Any) -> Tuple[Any, Dict[str, torch.Tensor], _SharedL self.pages.append(self._create_page(comm)) page = self.pages[-1] record = (index % self.records_per_page) * self.record_stride + if self.per_layer_sets: + self.sets += self._create_sets(1, comm) helper_set = self.sets[index % len(self.sets)] layer = _SharedLayer(self, index % len(self.sets)) self.layer_count += 1 @@ -886,13 +891,15 @@ def build_layer(self, comm: Any) -> Tuple[Any, Dict[str, torch.Tensor], _SharedL class SharedSlotArenaProvider: - """Provide one layer's arena over the two shared helper-slot banks. + """Provide one layer's arena over its helper-slot set. - The broadcaster requests the arena again after Parameters are bound, so a - provider memoizes its first build and requires identical geometry. + ``per_layer_helper_sets`` gives the layer its own helper set; otherwise it + maps one of the two banks shared by all layers. The broadcaster requests + the arena again after Parameters are bound, so a provider memoizes its + first build and requires identical geometry. """ - def __init__(self, owner: Any) -> None: + def __init__(self, owner: Any, *, per_layer_helper_sets: bool = False) -> None: # The model-level per-layer dynamic manager owns both the geometry cache # and creation order. No Mapping attribute or process-global registry # participates in allocation lifetime. @@ -903,6 +910,7 @@ def __init__(self, owner: Any) -> None: raise TypeError("SharedSlotArenaProvider owner must expose model-local pool state") self._pools = pools self._pool_order = order + self._per_layer_helper_sets = bool(per_layer_helper_sets) self.buffer_bytes = 0 self.group_sizes: Tuple[int, ...] = () self.pool: Optional[_SharedSlotPool] = None @@ -929,6 +937,7 @@ def build_hierarchical_live_arena( int(helper_count), int(device), bundle, + self._per_layer_helper_sets, ) if self._built is not None: if self._signature != signature: @@ -960,6 +969,7 @@ def build_hierarchical_live_arena( device=int(device), comm=comm, owner_registry=self._pool_order, + per_layer_sets=self._per_layer_helper_sets, ) self._pools[key] = pool self._built, self._tekit_views, self.layer, self.buffer_bytes = pool.build_layer(comm) diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py index dab7d320b36d..cad2bc82686b 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py @@ -79,6 +79,7 @@ def __init__( max_tokens_per_rank: int, copy_stream: torch.cuda.Stream, auxiliary_sms: int, + helper_residency: bool, layer_idx: Optional[int], ) -> None: from cuda.bindings import driver @@ -220,6 +221,7 @@ def __init__( tma_sm_count=self.auxiliary_sms, tma_warps=7, tma_route="plan", + helper_residency=bool(helper_residency), ) self.broadcaster.bind_scheduler(self.scheduler, self.scheduler_stream) @@ -449,8 +451,9 @@ def build_rebalance_slot_scheduler_group( ) -> RebalanceSlotSchedulerGroup: """Build one layer's scheduler/copy group from its manager's resources. - The manager supplies the model-owned EP communicator, copy stream and - auxiliary SM budget; the communicator is resolved before the stream. + The manager supplies the model-owned EP communicator, copy stream, + auxiliary SM budget and helper-slot mode; the communicator is resolved + before the stream. """ device = int(mapping.local_rank) return RebalanceSlotSchedulerGroup( @@ -464,6 +467,7 @@ def build_rebalance_slot_scheduler_group( device=device, copy_stream=manager.get_copy_stream(device), auxiliary_sms=int(manager.auxiliary_sms), + helper_residency=bool(manager.helper_residency), layer_idx=layer_idx, ) diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py index 29d8be77b92d..b0ae7de69549 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py @@ -390,6 +390,11 @@ def __init__( self.auxiliary_sms = int(getattr(config, "auxiliary_sms", 8)) if self.auxiliary_sms <= 0: raise ValueError("auxiliary_sms must be positive") + helper_slot_mode = getattr(config, "helper_slot_mode", "per_layer_resident") + if helper_slot_mode not in ("per_layer_resident", "shared_banks"): + raise ValueError(f"unknown helper_slot_mode {helper_slot_mode!r}") + # Resident helper experts need helper rows that only their layer writes. + self.helper_residency = helper_slot_mode == "per_layer_resident" # State read by the shared lifecycle and by the engine. self._previous_balancer = None @@ -489,7 +494,7 @@ def create_arena_provider(self) -> Any: """Create a per-layer provider scoped to this model manager.""" from .mega_moe.rebalance_live_arena import SharedSlotArenaProvider - provider = SharedSlotArenaProvider(self) + provider = SharedSlotArenaProvider(self, per_layer_helper_sets=self.helper_residency) self._arena_providers.append(provider) return provider diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py index d8dab8d7d622..cc74d5a6b34f 100644 --- a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py @@ -210,6 +210,7 @@ def __init__( tma_warps: int | None = None, tma_route: str = "plan", tma_source_load_factor: float = 1.0, + helper_residency: bool = False, ) -> None: for name, value in (("tma_sm_count", tma_sm_count), ("tma_warps", tma_warps)): if value is not None and (type(value) is not int or value <= 0): @@ -220,6 +221,11 @@ def __init__( raise ValueError("tma_warps must be at most 32; device shared-memory limits also apply") if tma_route not in ("plan", "source", "scatter"): raise ValueError("tma_route must be 'plan', 'source', or 'scatter'") + if type(helper_residency) is not bool: + raise ValueError("helper_residency must be a bool") + # Residency skips copies into helper rows that already hold the planned + # expert, so only this layer's copies may write those rows. + self.helper_residency = helper_residency if (type(tma_source_load_factor) not in (int, float) or not 0.01 <= tma_source_load_factor < 2.0 or abs(100 * tma_source_load_factor - round(100 * tma_source_load_factor)) > 1e-9): @@ -501,7 +507,13 @@ def bind_scheduler(self, scheduler: object, stream: object) -> None: try: bind(handle) torch.ops.trtllm.moe_rebalance_tma_bind_gpu_direct( - self.ctx, *tensors[1:], workspace, status, grid_sync, capacity + self.ctx, + *tensors[1:], + workspace, + status, + grid_sync, + capacity, + self.helper_residency, ) except BaseException: self._handoff_aborted = True diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index b3dcab68dd76..011743a8b92b 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1953,7 +1953,7 @@ class MoePerLayerDynamicLoadBalancerConfig(_MoeLoadBalancerConfigBase): status="prototype") num_slots: PositiveInt = Field( description= - "Total number of expert slots across all EP ranks. Must be divisible by `ep_size` and greater than the number of experts. Each rank keeps its `num_experts / ep_size` resident experts and adds `(num_slots - num_experts) / ep_size` helper slots. Unlike standard EPLB, the helper slots are not per-layer replica weights: they are two shared helper banks that all MoE layers alternate between and refill in every forward pass.", + "Total number of expert slots across all EP ranks. Must be divisible by `ep_size` and greater than the number of experts. Each rank keeps its `num_experts / ep_size` resident experts and adds `(num_slots - num_experts) / ep_size` helper slots. Unlike standard EPLB, the helper slots are not static replicas: every forward pass refills them from that layer's routing; see `helper_slot_mode` for their storage.", status="prototype") auxiliary_sms: int = Field( default=8, @@ -1962,6 +1962,11 @@ class MoePerLayerDynamicLoadBalancerConfig(_MoeLoadBalancerConfigBase): description= "Number of SMs that the routed-input quantization and shared-expert FC12 kernels leave free for the kernel that runs the HALO-Q scheduler and then the in-switch copy of its plan on the copy stream, one CTA per SM. Must be smaller than the device SM count.", status="prototype") + helper_slot_mode: Literal["per_layer_resident", "shared_banks"] = Field( + default="per_layer_resident", + description= + "Storage of the helper slots. `per_layer_resident` gives every MoE layer its own helper slots and skips copying an expert that a helper slot already holds from that layer's previous forward pass; it costs one helper set per MoE layer, about 6.4 GiB per GPU for DeepSeek-V4-Pro with 3 helper slots per rank. `shared_banks` has all MoE layers alternate between two shared helper banks and copy every helper expert in every forward pass.", + status="prototype") def validate_expert_capacity(self, num_experts: int) -> None: """Validate that every EP rank has at least one helper slot.""" diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index 6c0f63508985..2d2fca14a461 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -4141,6 +4141,17 @@ "title": "Auxiliary Sms", "type": "integer" }, + "helper_slot_mode": { + "default": "per_layer_resident", + "description": "Storage of the helper slots. `per_layer_resident` gives every MoE layer its own helper slots and skips copying an expert that a helper slot already holds from that layer's previous forward pass; it costs one helper set per MoE layer, about 6.4 GiB per GPU for DeepSeek-V4-Pro with 3 helper slots per rank. `shared_banks` has all MoE layers alternate between two shared helper banks and copy every helper expert in every forward pass.", + "enum": [ + "per_layer_resident", + "shared_banks" + ], + "status": "prototype", + "title": "Helper Slot Mode", + "type": "string" + }, "mode": { "const": "per_layer_dynamic", "default": "per_layer_dynamic", @@ -4150,7 +4161,7 @@ "type": "string" }, "num_slots": { - "description": "Total number of expert slots across all EP ranks. Must be divisible by `ep_size` and greater than the number of experts. Each rank keeps its `num_experts / ep_size` resident experts and adds `(num_slots - num_experts) / ep_size` helper slots. Unlike standard EPLB, the helper slots are not per-layer replica weights: they are two shared helper banks that all MoE layers alternate between and refill in every forward pass.", + "description": "Total number of expert slots across all EP ranks. Must be divisible by `ep_size` and greater than the number of experts. Each rank keeps its `num_experts / ep_size` resident experts and adds `(num_slots - num_experts) / ep_size` helper slots. Unlike standard EPLB, the helper slots are not static replicas: every forward pass refills them from that layer's routing; see `helper_slot_mode` for their storage.", "exclusiveMinimum": 0, "status": "prototype", "title": "Num Slots", diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index faffa16cef70..7ca5daa75c17 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -5705,6 +5705,17 @@ "title": "Auxiliary Sms", "type": "integer" }, + "helper_slot_mode": { + "default": "per_layer_resident", + "description": "Storage of the helper slots. `per_layer_resident` gives every MoE layer its own helper slots and skips copying an expert that a helper slot already holds from that layer's previous forward pass; it costs one helper set per MoE layer, about 6.4 GiB per GPU for DeepSeek-V4-Pro with 3 helper slots per rank. `shared_banks` has all MoE layers alternate between two shared helper banks and copy every helper expert in every forward pass.", + "enum": [ + "per_layer_resident", + "shared_banks" + ], + "status": "prototype", + "title": "Helper Slot Mode", + "type": "string" + }, "mode": { "const": "per_layer_dynamic", "default": "per_layer_dynamic", @@ -5714,7 +5725,7 @@ "type": "string" }, "num_slots": { - "description": "Total number of expert slots across all EP ranks. Must be divisible by `ep_size` and greater than the number of experts. Each rank keeps its `num_experts / ep_size` resident experts and adds `(num_slots - num_experts) / ep_size` helper slots. Unlike standard EPLB, the helper slots are not per-layer replica weights: they are two shared helper banks that all MoE layers alternate between and refill in every forward pass.", + "description": "Total number of expert slots across all EP ranks. Must be divisible by `ep_size` and greater than the number of experts. Each rank keeps its `num_experts / ep_size` resident experts and adds `(num_slots - num_experts) / ep_size` helper slots. Unlike standard EPLB, the helper slots are not static replicas: every forward pass refills them from that layer's routing; see `helper_slot_mode` for their storage.", "exclusiveMinimum": 0, "status": "prototype", "title": "Num Slots", diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 4e66e3e88fad..760c868f9112 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1076,6 +1076,15 @@ "kind": "value", "path": "moe_config.load_balancer.auxiliary_sms" }, + { + "allowed_values": [ + "per_layer_resident", + "shared_banks" + ], + "capture_policy": "literal", + "kind": "categorical", + "path": "moe_config.load_balancer.helper_slot_mode" + }, { "capture_policy": "int", "kind": "value", diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index 1053e5248844..600f89a82eb5 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -173,6 +173,54 @@ class Bundle: assert [pool.bundle for pool in FakePool.instances] == [first_bundle, second_bundle] +@pytest.mark.parametrize("per_layer_helper_sets", [False, True]) +def test_shared_slot_provider_pool_key_includes_helper_set_mode(per_layer_helper_sets): + source = _TORCH_ROOT / "moe/fused_moe/mega_moe/rebalance_live_arena.py" + + class FakePool: + instances = [] + + def __init__(self, **kwargs): + self.kwargs = kwargs + self._closed = False + self.group_sizes = () + self.instances.append(self) + + def build_layer(self, comm): + return object(), {}, object(), 0 + + provider_cls = _extract_class( + source, + "SharedSlotArenaProvider", + {"__init__", "build_hierarchical_live_arena"}, + { + "Any": Any, + "Dict": dict, + "Optional": Optional, + "Tuple": Tuple, + "_SharedSlotPool": FakePool, + }, + ) + owner = SimpleNamespace( + _rebalance_shared_slot_pools={}, + _rebalance_shared_slot_pool_order=[], + ) + kwargs = dict(world=8, rank=0, home_count=4, helper_count=1, bundle="bundle", device=0) + provider_cls(owner, per_layer_helper_sets=per_layer_helper_sets).build_hierarchical_live_arena( + comm=object(), **kwargs + ) + provider_cls( + owner, per_layer_helper_sets=not per_layer_helper_sets + ).build_hierarchical_live_arena(comm=object(), **kwargs) + + # One pool per mode: a per-layer pool never hands out a shared bank. + assert [pool.kwargs["per_layer_sets"] for pool in FakePool.instances] == [ + per_layer_helper_sets, + not per_layer_helper_sets, + ] + assert len(owner._rebalance_shared_slot_pools) == 2 + + def test_hierarchical_fused_submit_launches_bound_scheduler_and_copy(): source = _TORCH_ROOT / "moe/kernels/megamoe_scheduler/sami/hierarchical.py" launches = [] diff --git a/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py b/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py index 1259d2e131ea..a6a6b576d443 100644 --- a/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py +++ b/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py @@ -439,6 +439,28 @@ def create_arena_provider(): ) +@pytest.mark.parametrize( + ("helper_slot_mode", "per_layer_sets"), + [(None, True), ("per_layer_resident", True), ("shared_banks", False)], +) +def test_helper_slot_mode_selects_helper_residency(helper_slot_mode, per_layer_sets): + config = _Config() + if helper_slot_mode is not None: + config.helper_slot_mode = helper_slot_mode + manager = PerLayerDynamicMoeLoadBalancer(ep_rank=0, ep_size=4, config=config) + + assert manager.helper_residency is per_layer_sets + provider = manager.create_arena_provider() + assert provider._per_layer_helper_sets is per_layer_sets + + +def test_helper_slot_mode_rejects_unknown_mode(): + config = _Config() + config.helper_slot_mode = "private" + with pytest.raises(ValueError, match="helper_slot_mode"): + PerLayerDynamicMoeLoadBalancer(ep_rank=0, ep_size=4, config=config) + + def test_per_layer_dynamic_manager_shares_the_base_lifecycle_only(): manager = _manager() assert isinstance(manager, BaseMoeLoadBalancer) diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 2a38f5b712ec..5cecd7f6eee6 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -275,6 +275,7 @@ def test_resolves_auto_and_accepts_compatible_backend( MoePerLayerDynamicLoadBalancerConfig) assert config.backend == "MEGAMOE_CUTEDSL" assert config.load_balancer.auxiliary_sms == 8 + assert config.load_balancer.helper_slot_mode == "per_layer_resident" @pytest.mark.parametrize( "num_experts, message", @@ -341,6 +342,29 @@ def test_per_layer_dynamic_rejects_more_than_fused_copy_sm_limit( with pytest.raises(ValidationError, match="auxiliary_sms"): MoeConfig(load_balancer=values) + @pytest.mark.parametrize("mode", ["per_layer_resident", "shared_banks"]) + def test_per_layer_dynamic_accepts_helper_slot_modes(self, + mode: str) -> None: + values = self._per_layer_dynamic() + values["helper_slot_mode"] = mode + config = MoeConfig(load_balancer=values) + assert config.load_balancer.helper_slot_mode == mode + + def test_per_layer_dynamic_rejects_unknown_helper_slot_mode(self) -> None: + values = self._per_layer_dynamic() + values["helper_slot_mode"] = "per_layer" + with pytest.raises(ValidationError, match="helper_slot_mode"): + MoeConfig(load_balancer=values) + + def test_standard_mode_rejects_helper_slot_mode(self) -> None: + with pytest.raises(ValidationError, match="helper_slot_mode"): + MoeConfig(backend="CUTLASS", + load_balancer={ + "mode": "standard", + "num_slots": 64, + "helper_slot_mode": "shared_banks", + }) + def test_per_layer_dynamic_disables_decode_cuda_graph(self) -> None: args = TorchLlmArgs( model=llama_model_path, diff --git a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py index 5a76a0c0d105..b7859b540da6 100644 --- a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py +++ b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py @@ -210,11 +210,16 @@ def test_branch_snapshot_field_types( ({"num_slots": 288}, True), ({"mode": "standard", "num_slots": 288, "layer_updates_per_iter": 1}, True), ({"mode": "per_layer_dynamic", "num_slots": 320, "auxiliary_sms": 4}, True), + ( + {"mode": "per_layer_dynamic", "num_slots": 320, "helper_slot_mode": "shared_banks"}, + True, + ), # A mapping without mode is standard EPLB, which has no auxiliary_sms. ({"num_slots": 320, "auxiliary_sms": 4}, False), ({"mode": "per_layer_dynamic"}, False), ({"mode": "per_layer_dynamic", "num_slots": 320, "layer_updates_per_iter": 1}, False), ({"mode": "per_layer_dynamic", "num_slots": 320, "auxiliary_sms": 9}, False), + ({"mode": "per_layer_dynamic", "num_slots": 320, "helper_slot_mode": "private"}, False), ({"mode": "dynamic", "num_slots": 320}, False), ], ) From 9abf413b4390904422e83cd7f37c30bc41e81b1e Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 05:19:11 -0700 Subject: [PATCH 14/59] [None][test] Stub the raw stream switch in the fused-submit contract test The fused HALO-Q + TMA submit now switches to the copy stream through enter_cuda_stream/exit_cuda_stream. Provide both in the extracted-class namespace and check that the launch enters the copy stream and restores the caller's stream afterwards. Signed-off-by: jintaop --- .../_torch/moe/test_dynamic_eplb_contracts.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index 600f89a82eb5..a3b11a06bae1 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -229,9 +229,9 @@ def submit_fused(ctx, flags_mc, route_epoch, route_target, *scheduler_args): launches.append((ctx, flags_mc, route_epoch, route_target, scheduler_args)) torch = SimpleNamespace( - cuda=SimpleNamespace(device=lambda _: nullcontext(), stream=lambda _: nullcontext()), ops=SimpleNamespace(trtllm=SimpleNamespace(moe_rebalance_tma_submit_fused=submit_fused)), ) + stream_switches = [] ticket_cls = _extract_class( source, "HierarchicalCopyTicket", @@ -242,7 +242,12 @@ def submit_fused(ctx, flags_mc, route_epoch, route_target, *scheduler_args): source, "HierarchicalSamiWeightBroadcast", {"submit_fused", "_require_live_bank_reusable"}, - {"torch": torch, "HierarchicalCopyTicket": ticket_cls}, + { + "torch": torch, + "HierarchicalCopyTicket": ticket_cls, + "enter_cuda_stream": lambda stream: stream_switches.append(("enter", stream)) or stream, + "exit_cuda_stream": lambda state: stream_switches.append(("exit", state)), + }, ) outputs = object() requests = [] @@ -272,6 +277,9 @@ def fused_submit_args(routes, stream_handle): assert requests == [("routes", 19)] assert launches == [("ctx", 4096, "epoch", 8, ("scheduler", "routes"))] + # The launch runs on the copy stream and restores the caller's stream afterwards. + owner = broadcast._stream_owner + assert stream_switches == [("enter", owner), ("exit", owner)] assert (ticket.generation, ticket.command_count) == (1, 1) assert (broadcast._active_generation, broadcast._generation) == (1, 2) # The bank stays leased until the consumer of generation 1 is proven done. From dc83df43999b6bff4ed70f852d64c3b3ed5672e2 Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 06:00:33 -0700 Subject: [PATCH 15/59] [None][doc] Describe both helper-slot modes in the per-layer dynamic EPLB walkthrough Step 1 of the walkthrough still described the shared_banks layout as the only one. State the per_layer_resident default and the shared_banks alternative there, keep step 3 to the copy skip, and drop the word "shared" from the parallel-strategy summary. Signed-off-by: jintaop --- docs/source/features/parallel-strategy.md | 2 +- examples/wide_ep/ep_load_balancer/README.md | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/source/features/parallel-strategy.md b/docs/source/features/parallel-strategy.md index 8a39725f21b6..96ebb9a2198c 100644 --- a/docs/source/features/parallel-strategy.md +++ b/docs/source/features/parallel-strategy.md @@ -151,7 +151,7 @@ Wide-EP introduces the concept of **expert slots** that are decoupled from speci #### 3. Expert Parallelism Load Balancer (EPLB) - **Offline EPLB**: Pre-computed expert placement based on historical workload statistics - **Online EPLB**: Dynamic expert placement that adapts to real-time traffic patterns -- **Per-layer dynamic EPLB**: Expert placement re-planned for every MoE layer on every forward pass, with hot-expert weights copied into shared helper slots on the GPU (`MEGAMOE_CUTEDSL` backend only) +- **Per-layer dynamic EPLB**: Expert placement re-planned for every MoE layer on every forward pass, with hot-expert weights copied into helper slots on the GPU (`MEGAMOE_CUTEDSL` backend only) - Layer-wise weight redistribution to minimize inference disruption ### Architecture Overview diff --git a/examples/wide_ep/ep_load_balancer/README.md b/examples/wide_ep/ep_load_balancer/README.md index 850115bf8b0a..0ffb784680fa 100644 --- a/examples/wide_ep/ep_load_balancer/README.md +++ b/examples/wide_ep/ep_load_balancer/README.md @@ -276,9 +276,9 @@ The hottest rank is now at 70.0 against a mean of 65, so the imbalance drops fro Offline EPLB keeps one expert placement for the whole run, and online EPLB refreshes `layer_updates_per_iter` layers per iteration from decayed statistics. Per-layer dynamic EPLB instead re-plans every MoE layer on every forward pass, from the routing of that pass: -1. Each rank keeps its own `num_experts / ep_size` experts and has `(num_slots - num_experts) / ep_size` helper slots. The helper slots live in two banks that all MoE layers share, rather than in per-layer replicas. +1. Each rank keeps its own `num_experts / ep_size` experts and has `(num_slots - num_experts) / ep_size` helper slots. By default (`helper_slot_mode: per_layer_resident`) every MoE layer has its own helper slots; with `shared_banks`, all MoE layers alternate between two shared helper banks. 2. Before a layer's experts run, a GPU scheduler (HALO-Q) exchanges routed-token counts across the EP ranks, picks the hot experts to copy into helper slots, and maps every routed token to a resident or helper slot. -3. The same kernel then moves the picked expert weights over NVLink on a high-priority auxiliary stream, while the main stream runs other work such as input quantization and, for DeepSeek-V4, the shared experts. By default each layer keeps its own helper slots, and an expert that a helper slot already holds from the previous forward pass is not copied again. +3. The same kernel then moves the picked expert weights over NVLink on a high-priority auxiliary stream, while the main stream runs other work such as input quantization and, for DeepSeek-V4, the shared experts. By default, an expert that a helper slot already holds from that layer's previous forward pass is not copied again. 4. The MoE kernel reads a helper slot only after its weights are marked ready. Expert weights are not staged in host shared memory, so the GDRCopy and shared-memory setup that online EPLB needs does not apply. From 9bada715038741ca50c8158eb91b57021f7b5766 Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 06:01:49 -0700 Subject: [PATCH 16/59] [None][fix] Size the perfect-router expert-id table once and warn where it cannot apply ENABLE_PERFECT_ROUTER=expert_ids kept one table per expert geometry and grew it by replacing it with a larger one. That freed storage that CUDA graphs captured earlier still read, and a table rebuilt on the MoE chunking aux stream could be read on the main stream before it was written. The table is now sized once. MoE records model_config.max_num_tokens, the bound on the local tokens of every MoE call, in _init_perfect_router; the table is built with that many rows on the first call, not at construction, where the module can still be on the meta device. It is never replaced, and the current stream is synchronized after the build so other streams can read it. A call above the bound raises a ValueError, and a first call during CUDA graph capture raises a RuntimeError. Backends that route inside their kernel, such as TRTLLM without attention DP, take the scheduler's fused-routing branch, which has no expert ids to replace. That branch now logs a warning once when the mode is set, and the docs say where the mode applies and that layer-wise replay overrides it. Signed-off-by: jintaop --- docs/source/developer-guide/perf-analysis.md | 18 +++++-- .../_torch/moe/fused_moe/interface.py | 9 +++- .../_torch/moe/fused_moe/moe_scheduler.py | 11 ++++ tensorrt_llm/_torch/moe/fused_moe/routing.py | 48 ++++++++++++------ tests/unittest/_torch/moe/test_moe_routing.py | 50 +++++++++---------- 5 files changed, 86 insertions(+), 50 deletions(-) diff --git a/docs/source/developer-guide/perf-analysis.md b/docs/source/developer-guide/perf-analysis.md index 8e4830b51a32..37a66edc725d 100644 --- a/docs/source/developer-guide/perf-analysis.md +++ b/docs/source/developer-guide/perf-analysis.md @@ -248,11 +248,18 @@ the sender rank. The routing scales are left as computed. Use it where logits cannot steer the routing, for example the DeepSeek-V4 hash layers, which pick experts from the token IDs, or where the token count changes every iteration, as in context-only benchmarks: the IDs come from one cached -table per expert geometry, so no logits are synthesized per token count. +table per expert geometry, sized to `max_num_tokens`, so no logits are +synthesized per token count. -It applies to MoE backends that route through `ConfigurableMoE`, including -`MEGAMOE_CUTEDSL`. With the load balancer off, the gap to a plain run bounds -what any MoE load balancer can gain: +It applies where the MoE scheduler computes the routed expert IDs outside the +MoE kernel: always with `MEGAMOE_CUTEDSL`, `CUTLASS`, `DEEPGEMM` and `CUTEDSL`, +and with other backends under attention DP or with DeepSeek-V4 routing. A +backend that routes inside its kernel, such as `TRTLLM` without attention DP, +has no routed expert IDs to replace: the mode logs a warning and has no effect +there, so use `ENABLE_PERFECT_ROUTER=1` instead. + +With the load balancer off, the gap to a plain run bounds what any MoE load +balancer can gain: ```bash # Learned routing, load balancer off @@ -262,7 +269,8 @@ ENABLE_PERFECT_ROUTER=expert_ids trtllm-bench ... ``` The layer-wise benchmarks (`examples/layer_wise_benchmarks`) accept the same -environment variable with `--balance-method NotModified`. +environment variable with `--balance-method NotModified` and without +`--replay-file-path`, whose recorded routing would replace the balanced IDs. ### Supported Models diff --git a/tensorrt_llm/_torch/moe/fused_moe/interface.py b/tensorrt_llm/_torch/moe/fused_moe/interface.py index 6291e67d5a5e..07dea4102aed 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/interface.py +++ b/tensorrt_llm/_torch/moe/fused_moe/interface.py @@ -392,7 +392,7 @@ def __init__( # un-overridden default would be ``expert_size_per_partition = # num_experts`` (each rank allocating storage for every expert). self._init_dwdp_expert_layout() - self._init_perfect_router() + self._init_perfect_router(model_config) def _init_dwdp_expert_layout(self): """Override expert layout when DWDP is enabled. @@ -466,13 +466,17 @@ def _get_perfect_router_dtype(self) -> torch.dtype: return torch.float32 return self.dtype if self.dtype is not None else torch.float32 - def _init_perfect_router(self): + def _init_perfect_router(self, model_config: ModelConfig): # ENABLE_PERFECT_ROUTER=1 replaces the router logits before routing; # ENABLE_PERFECT_ROUTER=expert_ids replaces the routed expert ids after # routing. Both are for performance analysis only. perfect_router = os.environ.get("ENABLE_PERFECT_ROUTER", "0") self._enable_perfect_router = perfect_router == "1" self._enable_perfect_router_expert_ids = perfect_router == "expert_ids" + # No MoE call routes more local tokens than this. It sizes the + # expert-id table, which is built on the first call, not here: the + # module may still be on the meta device. + self._perfect_router_max_num_tokens = model_config.max_num_tokens if not self._enable_perfect_router: return @@ -520,6 +524,7 @@ def _maybe_get_perfect_router_expert_ids( experts_per_token=experts_per_token, moe_ep_size=self.ep_size, ep_rank=self.ep_rank, + max_num_tokens=self._perfect_router_max_num_tokens, device=token_selected_experts.device) def _init_load_balancer( diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py index 2076be15b560..c5708c1cdac8 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py @@ -51,6 +51,7 @@ from tensorrt_llm._torch.moe.expert_statistic import ExpertStatistic from tensorrt_llm._torch.route_capture import get_active_route_capture # R3 from tensorrt_llm._torch.utils import EventType, Fp4QuantizedTensor +from tensorrt_llm.logger import logger from tensorrt_llm.tools.layer_wise_benchmarks import get_calibrator from .communication import DeepEP, DeepEPLowLatency, NcclEP, NVLinkOneSided, NVLinkTwoSided @@ -543,6 +544,16 @@ def _forward_chunk_impl( f"EPLB is enabled but backend {moe.backend.__class__.__name__} " f"has fused routing (does not support routing separation)" ) + if moe._enable_perfect_router_expert_ids: + logger.warning_once( + "ENABLE_PERFECT_ROUTER=expert_ids has no effect on " + f"{moe.backend.__class__.__name__} here: it routes inside its " + "kernel, so there are no routed expert ids to replace. Use " + "ENABLE_PERFECT_ROUTER=1 instead; for TRTLLM, setting " + "TLLM_TRTLLMGEN_FORCE_SEPARATED_ROUTING=1 in every compared run " + "also works.", + key="perfect_router_expert_ids_fused_routing", + ) token_selected_experts = None token_final_scales = None diff --git a/tensorrt_llm/_torch/moe/fused_moe/routing.py b/tensorrt_llm/_torch/moe/fused_moe/routing.py index f9b9c86201ce..7e9823b57122 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/routing.py +++ b/tensorrt_llm/_torch/moe/fused_moe/routing.py @@ -16,11 +16,8 @@ # Global cache for perfect router logits shared across MoE modules _PERFECT_ROUTER_LOGITS_CACHE = {} # Global cache for perfect router expert-id tables shared across MoE modules, -# keyed by expert geometry, sender rank and device +# keyed by expert geometry, sender rank, token bound and device _PERFECT_ROUTER_EXPERT_IDS_CACHE = {} -# Initial expert-id table rows, enough for the per-rank token counts of typical -# prefill chunks, so the table is normally built once, during warmup -_PERFECT_ROUTER_EXPERT_IDS_MIN_ROWS = 16384 def _get_perfect_router_cache_signature(routing_method): @@ -243,6 +240,7 @@ def get_cached_perfect_router_logits(num_tokens: int, num_experts: int, def get_cached_perfect_router_expert_ids(num_tokens: int, num_experts: int, experts_per_token: int, moe_ep_size: int, ep_rank: int, + max_num_tokens: int, device: torch.device) -> torch.Tensor: """ Return perfectly balanced expert ids for ``num_tokens`` local tokens. @@ -257,10 +255,11 @@ def get_cached_perfect_router_expert_ids(num_tokens: int, num_experts: int, slot ``t * experts_per_token + j``, rotated by ``ep_rank``, goes to EP rank ``slot % moe_ep_size`` and to expert ``(slot // moe_ep_size) % (num_experts // moe_ep_size)`` on it. A slot's expert depends only on its - own index, so the ids for ``n`` tokens are the leading ``n`` rows of any - longer table. One table per geometry and device is built on first use and - grown by doubling; every call returns a view of its leading rows, with no - kernel launch. + own index, so the ids for ``n`` tokens are the leading ``n`` rows of a + longer table. One table of ``max_num_tokens`` rows per geometry, bound and + device is built on the first call and never replaced, so CUDA graphs + captured later keep reading valid memory. Every call returns a view of its + leading rows, with no kernel launch. Args: num_tokens: Number of local tokens on the current EP rank. @@ -268,22 +267,39 @@ def get_cached_perfect_router_expert_ids(num_tokens: int, num_experts: int, experts_per_token: Number of experts selected for each token. moe_ep_size: Number of EP ranks. ep_rank: Sender EP rank used to rotate the assignment stream. + max_num_tokens: Upper bound on ``num_tokens`` in every call; it sizes + the table. device: Device of the routed expert ids. Returns: torch.Tensor: Contiguous int32 expert ids of shape ``[num_tokens, experts_per_token]``. + + Raises: + ValueError: If ``num_tokens`` exceeds ``max_num_tokens``. + RuntimeError: If the table would be built while a CUDA graph is being + captured. """ - cache_key = (num_experts, experts_per_token, moe_ep_size, ep_rank, device) + if num_tokens > max_num_tokens: + raise ValueError( + f"ENABLE_PERFECT_ROUTER=expert_ids got {num_tokens} tokens in one " + f"MoE call, more than max_num_tokens ({max_num_tokens}), which " + "sizes its expert-id table") + cache_key = (num_experts, experts_per_token, moe_ep_size, ep_rank, + max_num_tokens, device) table = _PERFECT_ROUTER_EXPERT_IDS_CACHE.get(cache_key) - if table is None or table.shape[0] < num_tokens: - if table is None: - num_rows = _PERFECT_ROUTER_EXPERT_IDS_MIN_ROWS - else: - num_rows = 2 * table.shape[0] + if table is None: + if device.type == "cuda" and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "ENABLE_PERFECT_ROUTER=expert_ids builds its expert-id table on " + "the first MoE call, which must run before CUDA graph capture") table = FlatTopKPerfectRouterPlanner._get_balanced_expert_indices( - max(num_rows, num_tokens), num_experts, experts_per_token, - moe_ep_size, ep_rank, device).to(torch.int32) + max_num_tokens, num_experts, experts_per_token, moe_ep_size, + ep_rank, device).to(torch.int32) + if table.is_cuda: + # Other streams, such as the MoE chunking aux stream, read the + # table without waiting for the stream that built it. + torch.cuda.current_stream(table.device).synchronize() _PERFECT_ROUTER_EXPERT_IDS_CACHE[cache_key] = table return table[:num_tokens] diff --git a/tests/unittest/_torch/moe/test_moe_routing.py b/tests/unittest/_torch/moe/test_moe_routing.py index 1c286264cffd..babdd4d24e60 100644 --- a/tests/unittest/_torch/moe/test_moe_routing.py +++ b/tests/unittest/_torch/moe/test_moe_routing.py @@ -650,25 +650,34 @@ def test_perfect_router_expert_ids_balanced(num_experts, top_k, moe_ep_size): - every EP rank and every expert receives an equal share of the routes, within one; - a token's routes land on distinct EP ranks and distinct experts; - - the ids for fewer tokens are a view of the leading rows of the cached - table, also after the table grows. + - every call returns a view of the same table, sized to the token bound; + - a call above the bound raises instead of growing the table. """ device = torch.device("cpu") experts_per_rank = num_experts // moe_ep_size + max_num_tokens = 8192 moe_routing._PERFECT_ROUTER_EXPERT_IDS_CACHE.clear() + def get_ids(num_tokens, ep_rank): + return get_cached_perfect_router_expert_ids( + num_tokens=num_tokens, + num_experts=num_experts, + experts_per_token=top_k, + moe_ep_size=moe_ep_size, + ep_rank=ep_rank, + max_num_tokens=max_num_tokens, + device=device) + for ep_rank in range(moe_ep_size): - largest = None - for num_tokens in [1, 3, 7, 100, 1000, 8191]: - ids = get_cached_perfect_router_expert_ids(num_tokens=num_tokens, - num_experts=num_experts, - experts_per_token=top_k, - moe_ep_size=moe_ep_size, - ep_rank=ep_rank, - device=device) + table_ptr = None + for num_tokens in [1, 3, 7, 100, 1000, max_num_tokens]: + ids = get_ids(num_tokens, ep_rank) assert ids.shape == (num_tokens, top_k) assert ids.dtype == torch.int32 assert ids.is_contiguous() + if table_ptr is None: + table_ptr = ids.data_ptr() + assert ids.data_ptr() == table_ptr rank_counts = torch.bincount(ids.flatten() // experts_per_rank, minlength=moe_ep_size) @@ -680,24 +689,11 @@ def test_perfect_router_expert_ids_balanced(num_experts, top_k, moe_ep_size): ranks = ids // experts_per_rank assert all(len(set(row)) == top_k for row in ranks.tolist()) assert all(len(set(row)) == top_k for row in ids.tolist()) - largest = ids - # Growing the table keeps the leading rows, so cached views stay valid. - grown = get_cached_perfect_router_expert_ids( - num_tokens=2 * moe_routing._PERFECT_ROUTER_EXPERT_IDS_MIN_ROWS + 1, - num_experts=num_experts, - experts_per_token=top_k, - moe_ep_size=moe_ep_size, - ep_rank=ep_rank, - device=device) - assert torch.equal(grown[:largest.shape[0]], largest) - again = get_cached_perfect_router_expert_ids(num_tokens=5, - num_experts=num_experts, - experts_per_token=top_k, - moe_ep_size=moe_ep_size, - ep_rank=ep_rank, - device=device) - assert again.data_ptr() == grown.data_ptr() + with pytest.raises(ValueError, match="max_num_tokens"): + get_ids(max_num_tokens + 1, ep_rank) + + assert len(moe_routing._PERFECT_ROUTER_EXPERT_IDS_CACHE) == moe_ep_size # ----------------------------------------------------------------- From ed539ce0a60f1f3868ff2cce301b4d2c5616df52 Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 06:03:22 -0700 Subject: [PATCH 17/59] [None][fix] Layer-wise benchmarks: reject unsupported load-balancer and replay setups - run.py rejects --replay-file-path together with ENABLE_PERFECT_ROUTER=expert_ids. The calibrator replays the recorded expert IDs after the substitution, so the perfect-router arm silently measured the recorded routing. - Runner rejects, before the model loads, per-layer dynamic EPLB with CUDA graphs, which failed at warmup with a scheduler error, and online standard EPLB, which synchronizes every MoE layer of the model while only a slice runs here and failed in finalize_model or hung. - README: the DeepSeek-V4 commands pass --tokens-per-block 128, which its cache manager requires, and --max-seq-len 512, since 256 or less trips an assert in its sparse-attention top-k. They use the tested 64 x 128 batch and the FP4 combine. The text now says that replay overrides the expert_ids substitution and that a recording fixes num_slots, so one recording cannot serve both the load-balancer-off and per-layer dynamic arms. Signed-off-by: jintaop --- examples/layer_wise_benchmarks/README.md | 21 +++++++++++++------ examples/layer_wise_benchmarks/run.py | 9 +++++++- .../tools/layer_wise_benchmarks/runner.py | 16 ++++++++++++++ 3 files changed, 39 insertions(+), 7 deletions(-) diff --git a/examples/layer_wise_benchmarks/README.md b/examples/layer_wise_benchmarks/README.md index 7bb0668f6164..c9c70d311377 100644 --- a/examples/layer_wise_benchmarks/README.md +++ b/examples/layer_wise_benchmarks/README.md @@ -125,15 +125,24 @@ NP=16 ./mpi_launch.sh ./run.sh config_gen.yaml --model $K3 --layer-indices 4,5,6 # load balancer off, per-layer dynamic EPLB (num_slots = experts + helper slots per rank * EP, # here 384 + 3 * 8), and perfectly balanced routing with the load balancer off (the upper # bound for any MoE load balancer). MEGAMOE_CUTEDSL is not in the balance-method whitelist, -# so the arms use --balance-method NotModified. To replay the routing of an end-to-end run, -# add --replay-file-path (see "Performance alignment" below); the third arm replaces the -# routed expert IDs either way. A slice may start mid-model (here 7,8). -NP=8 ./mpi_launch.sh ./run.sh config_ctx.yaml --model $DSV4 --layer-indices 7,8 --moe-backend MEGAMOE_CUTEDSL --balance-method NotModified -NP=8 ./mpi_launch.sh ./run.sh config_ctx.yaml --model $DSV4 --layer-indices 7,8 --moe-backend MEGAMOE_CUTEDSL --balance-method NotModified --moe-load-balancer eplb_per_layer_dynamic.yaml -NP=8 ./mpi_launch.sh -x ENABLE_PERFECT_ROUTER=expert_ids ./run.sh config_ctx.yaml --model $DSV4 --layer-indices 7,8 --moe-backend MEGAMOE_CUTEDSL --balance-method NotModified +# so the arms use --balance-method NotModified. DeepSeek-V4 needs --tokens-per-block 128 or +# 256, and --max-seq-len above 256 (an assert in its sparse-attention top-k); 512 is tested. +# MEGAMOE_COMBINE_FORMAT=16e2m1xbf16 selects the FP4 combine. A slice may start mid-model +# (here 7,8). +NP=8 ./mpi_launch.sh -x MEGAMOE_COMBINE_FORMAT=16e2m1xbf16 ./run.sh config_ctx.yaml --model $DSV4 --layer-indices 7,8 --tokens-per-block 128 --batch-size 64 --seq-len-q 128 --max-seq-len 512 --moe-backend MEGAMOE_CUTEDSL --balance-method NotModified +NP=8 ./mpi_launch.sh -x MEGAMOE_COMBINE_FORMAT=16e2m1xbf16 ./run.sh config_ctx.yaml --model $DSV4 --layer-indices 7,8 --tokens-per-block 128 --batch-size 64 --seq-len-q 128 --max-seq-len 512 --moe-backend MEGAMOE_CUTEDSL --balance-method NotModified --moe-load-balancer eplb_per_layer_dynamic.yaml +NP=8 ./mpi_launch.sh -x MEGAMOE_COMBINE_FORMAT=16e2m1xbf16 -x ENABLE_PERFECT_ROUTER=expert_ids ./run.sh config_ctx.yaml --model $DSV4 --layer-indices 7,8 --tokens-per-block 128 --batch-size 64 --seq-len-q 128 --max-seq-len 512 --moe-backend MEGAMOE_CUTEDSL --balance-method NotModified # with eplb_per_layer_dynamic.yaml holding: # mode: per_layer_dynamic # num_slots: 408 +# Per-layer dynamic EPLB runs without CUDA graphs only, and online standard EPLB is not +# supported on a layer slice; the runner rejects both. +# To replay the routing of an end-to-end run, add --replay-file-path (see "Performance +# alignment" below) to the first two arms only, each with a recording made with the same +# num_slots: a recording with the load balancer off (384 slots) cannot be replayed into the +# per-layer dynamic arm (408 slots), nor the reverse. Replay overrides the expert IDs that +# ENABLE_PERFECT_ROUTER=expert_ids substitutes, so run the third arm without it; run.py +# rejects that combination. Replay with MEGAMOE_CUTEDSL and the CTX phase is untested. # Run with DeepEP A2A NP=4 ./mpi_launch.sh -x TRTLLM_FORCE_COMM_METHOD=DEEPEP ./run.sh config_ctx.yaml --moe-backend CUTEDSL diff --git a/examples/layer_wise_benchmarks/run.py b/examples/layer_wise_benchmarks/run.py index 476c6ce18114..66454b30d1f3 100644 --- a/examples/layer_wise_benchmarks/run.py +++ b/examples/layer_wise_benchmarks/run.py @@ -80,7 +80,9 @@ def comma_separated_floats(s): "--moe-load-balancer", type=str, help="MoeConfig.load_balancer: a YAML file path, or a mapping when set in the config" - " file, e.g. {mode: per_layer_dynamic, num_slots: 408}. Default: no load balancer.", + " file, e.g. {mode: per_layer_dynamic, num_slots: 408}. Supports per_layer_dynamic" + " without CUDA graphs, and static standard EPLB (layer_updates_per_iter 0)." + " Default: no load balancer.", ) parser.add_argument( "--spec-max-draft-len", @@ -185,6 +187,11 @@ def comma_separated_floats(s): args.use_cuda_graph = False if (args.replay_start_iter is None) != (args.replay_stop_iter is None): parser.error("Both --replay-start-iter and --replay-stop-iter must be provided or none") +if args.replay_file_path and os.environ.get("ENABLE_PERFECT_ROUTER") == "expert_ids": + parser.error( + "--replay-file-path replays the recorded expert IDs over the ones that" + " ENABLE_PERFECT_ROUTER=expert_ids substitutes; run the perfect-router arm without it" + ) if args.replay_verify_metadata is None: args.replay_verify_metadata = True if args.spec_max_draft_len is not None: diff --git a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py index 32ffb5893e1f..c48a96278a2a 100644 --- a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py +++ b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py @@ -451,6 +451,22 @@ def __init__( dtype=kv_cache_dtype, mamba_ssm_cache_dtype=mamba_ssm_cache_dtype ), ) + # Reject load balancers that cannot run on a layer slice driven by this + # harness. The validator above has resolved a YAML path to its config. + load_balancer_config = llm_args.moe_config.load_balancer + load_balancer_mode = getattr(load_balancer_config, "mode", None) + if load_balancer_mode == "per_layer_dynamic" and use_cuda_graph: + raise ValueError( + "Per-layer dynamic EPLB does not support CUDA graphs: it binds one submission" + " stream and keeps host-side plan state that a graph replay cannot advance." + " Rerun with --no-use-cuda-graph." + ) + if load_balancer_mode == "standard" and load_balancer_config.layer_updates_per_iter > 0: + raise ValueError( + "Layer-wise benchmarks support per_layer_dynamic EPLB and static standard EPLB" + " (layer_updates_per_iter 0) only: online standard EPLB synchronizes every MoE" + " layer of the model, but only a slice of the layers runs here." + ) with ( self.scaled_from_ctx(scaled_from, mapping), self.vision_config_ctx(vision_config), From e88384d5c197f9052c060f4cf7d997f3ad0e2e23 Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 08:01:11 -0700 Subject: [PATCH 18/59] [None][fix] Give the HALO-Q scheduler the MoE communication warmup timeout during warmup HALO-Q waits for every EP rank inside the kernel, in its count exchange and, in the fused kernel, for the route epoch, and traps once a peer is later than SPIN_CYCLES (8e9 cycles, about 4 s). Executor warmup and AutoTuner profiling skew healthy ranks by far more than that: on VR200 the ON arm died in warmup with CUDA error 719 inside the plan-gap shared-expert FC12 profiling at num_tokens=8192, which then surfaced as a failed cuStreamWaitValue64. The scheduler already takes its spin bound as a kernel argument, so choose it per launch: - SchedulerSpinBound is a MoE communication timeout proxy, one per model, registered by PerLayerDynamicMoeLoadBalancer and released at shutdown. The model engine's is_warmup transitions switch it like every other MoE communication kernel: the warmup timeout (1800 s by default) during warmup and the serving timeout or the native SPIN_CYCLES otherwise. Launches while AutoTuner is tuning, such as profiling outside the executor's warmup, also get the warmup timeout. - fused_submit_args, submit_fused and the standalone launch/submit/run take spin_cycles; the startup HALO-Q launch always uses the warmup bound. - Timeouts convert to cycles at the nominal 2 GHz clock that the MoE all-to-all kernels assume. No kernel or C++ change; no new environment variable. Signed-off-by: jintaop --- .../moe/fused_moe/MOE_DEVELOPER_GUIDE.md | 5 + .../mega_moe/rebalance_slot_scheduler.py | 65 ++++- .../moe/fused_moe/per_layer_dynamic_eplb.py | 28 +++ .../cuda_scheduler/runtime.py | 63 ++++- .../megamoe_scheduler/sami/hierarchical.py | 15 +- .../_torch/moe/test_dynamic_eplb_contracts.py | 227 +++++++++++++++++- .../_torch/moe/test_per_layer_dynamic_eplb.py | 21 ++ 7 files changed, 398 insertions(+), 26 deletions(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md index 5e3eb29e1185..8b38336449c7 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -879,6 +879,11 @@ set wrong has a test that catches it: memory wait on that epoch orders it after the routes. A single CPU submitter owns both streams; the path does not require a CUDA green context or a steady-state host wait. +- HALO-Q waits for every EP rank inside the kernel and traps once a peer is + later than its spin bound. Each launch passes the bound, which + `SchedulerSpinBound` takes from the MoE communication timeout guard: the + warmup timeout during warmup and AutoTuner profiling, the native bound + otherwise. - `helper_slot_mode=per_layer_resident` gives each layer its own helper-slot set, and the copy keeps a per-layer table of the expert each helper slot holds so that it skips experts a whole team already holds. Any other writer diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py index cad2bc82686b..dac4279e21f9 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py @@ -21,13 +21,57 @@ from __future__ import annotations import threading -from typing import Any, Optional, Tuple +from typing import Any, Callable, Optional, Tuple import torch from tensorrt_llm.logger import logger -__all__ = ["RebalanceSlotSchedulerGroup", "build_rebalance_slot_scheduler_group"] +__all__ = [ + "RebalanceSlotSchedulerGroup", + "SchedulerSpinBound", + "build_rebalance_slot_scheduler_group", +] + + +class SchedulerSpinBound: + """Bound of the HALO-Q scheduler's cross-rank waits, chosen per launch. + + HALO-Q exchanges route counts with every EP rank inside the kernel and + traps once a peer is later than this bound. During warmup, one-time JIT + compilation and autotuning can delay a healthy rank by minutes. As a MoE + communication timeout proxy, the bound follows every warmup transition: + warmup gets the warmup timeout, and serving the serving timeout or the + scheduler's native ``SPIN_CYCLES``. Launches while ``is_tuning`` returns + True, such as AutoTuner profiling outside warmup, also get the warmup + timeout. One bound serves every layer of a model. + + Args: + warmup_seconds: The MoE communication warmup timeout. + is_tuning: Returns whether kernels are being profiled. + """ + + name = "per-layer dynamic EPLB HALO-Q" + + def __init__(self, *, warmup_seconds: int, is_tuning: Callable[[], bool]) -> None: + from tensorrt_llm._torch.moe.kernels.megamoe_scheduler.cuda_scheduler.runtime import ( + spin_cycles_for_timeout, + ) + + self._spin_cycles_for_timeout = spin_cycles_for_timeout + self._is_tuning = is_tuning + self.warmup_cycles = spin_cycles_for_timeout(warmup_seconds) + self._timeout_cycles = spin_cycles_for_timeout(None) + + def set_timeout_seconds(self, seconds: Optional[int]) -> None: + """Apply a timeout to the launches issued after this call; ``None`` restores the default.""" + self._timeout_cycles = self._spin_cycles_for_timeout(seconds) + + def cycles(self) -> int: + """Return the bound of the next launch.""" + if self._is_tuning(): + return self.warmup_cycles + return self._timeout_cycles class _LiveBankLeaseProvider: @@ -80,6 +124,7 @@ def __init__( copy_stream: torch.cuda.Stream, auxiliary_sms: int, helper_residency: bool, + spin_bound: SchedulerSpinBound, layer_idx: Optional[int], ) -> None: from cuda.bindings import driver @@ -109,6 +154,7 @@ def __init__( self.topk = int(topk) self.max_tokens_per_rank = int(max_tokens_per_rank) self.auxiliary_sms = int(auxiliary_sms) + self._spin_bound = spin_bound if copy_stream.device.index != self.device: raise ValueError("copy stream device does not match scheduler device") self.arena = arena @@ -204,7 +250,9 @@ def __init__( # Warm all ranks before binding the plan channel: a bound launch would # leave a pending publication that must be consumed by the broadcaster. # The scheduler input is already initialized and peer allocations synchronized. - self.scheduler.launch(self.scheduler_stream) + # Loading the kernel can still skew the ranks, so this launch always + # waits for its peers as long as warmup does. + self.scheduler.launch(self.scheduler_stream, spin_cycles=spin_bound.warmup_cycles) torch.cuda.synchronize(self.device) ep_comm.barrier() @@ -374,7 +422,11 @@ def plan_schedule(self, logical_expert_ids: torch.Tensor) -> Tuple[Any, int]: self._pending_release = None route_target = self._route_target + self.auxiliary_sms ticket = self.broadcaster.submit_fused( - outputs, logical_expert_ids, self._route_epoch, route_target + outputs, + logical_expert_ids, + self._route_epoch, + route_target, + spin_cycles=self._spin_bound.cycles(), ) if int(ticket.generation) != self._scheduled_generation: raise RuntimeError("TMA and HALO-Q generation mismatch") @@ -452,8 +504,8 @@ def build_rebalance_slot_scheduler_group( """Build one layer's scheduler/copy group from its manager's resources. The manager supplies the model-owned EP communicator, copy stream, - auxiliary SM budget and helper-slot mode; the communicator is resolved - before the stream. + auxiliary SM budget, helper-slot mode and scheduler spin bound; the + communicator is resolved before the stream. """ device = int(mapping.local_rank) return RebalanceSlotSchedulerGroup( @@ -468,6 +520,7 @@ def build_rebalance_slot_scheduler_group( copy_stream=manager.get_copy_stream(device), auxiliary_sms=int(manager.auxiliary_sms), helper_residency=bool(manager.helper_residency), + spin_bound=manager.get_scheduler_spin_bound(), layer_idx=layer_idx, ) diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py index b0ae7de69549..f8fe733bc653 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py @@ -407,6 +407,7 @@ def __init__( self._copy_stream: Optional[torch.cuda.Stream] = None self._copy_stream_device: Optional[int] = None self._ep_comm: Optional[Any] = None + self._scheduler_spin_bound: Optional[Any] = None # Providers use this manager as their scope owner. Resource references # stay reachable here, never through Mapping or a process-global registry. @@ -490,6 +491,28 @@ def get_ep_comm(self, process_group: Any) -> Any: self._ep_comm.validate_geometry(expected_rank=self.ep_rank, expected_world=self.ep_size) return self._ep_comm + def get_scheduler_spin_bound(self) -> Any: + """Return the HALO-Q spin bound shared by this model's layers. + + It is registered with the MoE communication timeout guard, which + lengthens it while the engine warms up. + """ + if self._scheduler_spin_bound is None: + from ...autotuner import AutoTuner + from .mega_moe.rebalance_slot_scheduler import SchedulerSpinBound + from .moe_comm_timeout_guard import ( + get_moe_comm_timeout_budgets, + register_moe_comm_timeout_proxy, + ) + + spin_bound = SchedulerSpinBound( + warmup_seconds=get_moe_comm_timeout_budgets().warmup_seconds, + is_tuning=lambda: AutoTuner.get().is_tuning_mode, + ) + register_moe_comm_timeout_proxy(spin_bound) + self._scheduler_spin_bound = spin_bound + return self._scheduler_spin_bound + def create_arena_provider(self) -> Any: """Create a per-layer provider scoped to this model manager.""" from .mega_moe.rebalance_live_arena import SharedSlotArenaProvider @@ -596,6 +619,11 @@ def shutdown(self, local_safe: bool = True) -> bool: layer._clear_runtime() self._arena_providers.clear() self._ep_comm = None + if self._scheduler_spin_bound is not None: + from .moe_comm_timeout_guard import unregister_moe_comm_timeout_proxy + + unregister_moe_comm_timeout_proxy(self._scheduler_spin_bound) + self._scheduler_spin_bound = None self._rebalance_shared_slot_pools.clear() self._rebalance_shared_slot_pool_order.clear() self._copy_stream = None diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py index d9dd5b328e98..0bd68b7fb0fe 100644 --- a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py @@ -16,7 +16,27 @@ MAX_EXPERTS = 384 MAX_CTAS = 128 MAX_BROADCASTS = MAX_EXPERTS +# Default bound, in SM clock cycles, of the scheduler's grid barrier and +# cross-rank waits; a timeout records a status flag. Each launch takes its bound +# as an argument, so callers can allow more peer skew while ranks warm up. SPIN_CYCLES = 8_000_000_000 +# Nominal clock64 rate that converts a timeout in seconds to SM cycles, as the +# MoE all-to-all kernels assume. SPIN_CYCLES is about 4 s at this rate. +ASSUMED_CLOCK_HZ = 2_000_000_000 + + +def spin_cycles_for_timeout(seconds: int | None) -> int: + """Spin bound for a timeout in seconds; ``None`` selects ``SPIN_CYCLES``.""" + if seconds is None: + return SPIN_CYCLES + if type(seconds) is not int or seconds <= 0: + raise ValueError("scheduler timeout must be a positive exact int of seconds") + return seconds * ASSUMED_CLOCK_HZ + + +def _check_spin_cycles(spin_cycles: int) -> None: + if type(spin_cycles) is not int or spin_cycles <= 0: + raise ValueError("spin_cycles must be a positive exact int") def _plan_channel_word_count(*, abi_version: int, channel_words: int | None, route_features: int, ep: int, helpers: int) -> int: @@ -442,7 +462,10 @@ def release_plan_channel(self) -> None: raise RuntimeError("plan channel has no pending publication") self._plan_channel_pending = False - def launch(self, stream: object | None = None) -> CudaSchedulerOutputs: + def launch( + self, stream: object | None = None, *, spin_cycles: int = SPIN_CYCLES + ) -> CudaSchedulerOutputs: + _check_spin_cycles(spin_cycles) if not self._connected: raise RuntimeError("EP peer bases are not connected") with self._plan_channel_lock: @@ -467,6 +490,7 @@ def launch(self, stream: object | None = None) -> CudaSchedulerOutputs: self.logical_expert_ids, int(handle), valid_route_count=self.cfg.route_count, + spin_cycles=spin_cycles, ) if self._plan_channel_bound: self._plan_channel_pending = True @@ -478,6 +502,7 @@ def _launch_op( stream_handle: int, *, valid_route_count: int, + spin_cycles: int, ) -> None: stream = self._plan_channel_torch_stream if stream is None or _stream_handle(stream) != stream_handle: @@ -485,13 +510,18 @@ def _launch_op( state = enter_cuda_stream(stream) try: torch.ops.trtllm.moe_rebalance_halo_q( - *self._halo_q_args(routes, valid_route_count, enable_pdl=self.cfg.enable_pdl) + *self._halo_q_args( + routes, + valid_route_count, + enable_pdl=self.cfg.enable_pdl, + spin_cycles=spin_cycles, + ) ) finally: exit_cuda_stream(state) def _halo_q_args( - self, routes: torch.Tensor, valid_route_count: int, *, enable_pdl: bool + self, routes: torch.Tensor, valid_route_count: int, *, enable_pdl: bool, spin_cycles: int ) -> tuple: """Positional moe_rebalance_halo_q arguments, shared by the fused copy op.""" return ( @@ -516,7 +546,7 @@ def _halo_q_args( self.cfg.ctas, self.cfg.threads, enable_pdl, - SPIN_CYCLES, + spin_cycles, self._plan_channel_abi_version, self._plan_channel_words, self._plan_channel_route_features, @@ -526,6 +556,7 @@ def _halo_q_args( def submit( self, routes: torch.Tensor, stream_handle: int, *, valid_tokens: int | None = None, + spin_cycles: int = SPIN_CYCLES, ) -> CudaSchedulerOutputs: """Read a caller-owned CUDA int32[T,K] tensor without input staging. @@ -545,6 +576,7 @@ def submit( with changing pointers/lengths use eager submit; stage()+launch() remains available for a fixed-capacity graph input buffer. """ + _check_spin_cycles(spin_cycles) with self._plan_channel_lock: if self._plan_channel_bound: if self._plan_channel_pending: @@ -558,18 +590,24 @@ def submit( "CUDA scheduler uses one static stream because its workspace " "and reusable barrier epochs are shared" ) - self._submit_routes(routes, stream_handle, valid_tokens=valid_tokens) + self._submit_routes( + routes, stream_handle, valid_tokens=valid_tokens, spin_cycles=spin_cycles + ) if self._plan_channel_bound: self._plan_channel_pending = True return self.outputs - def fused_submit_args(self, routes: torch.Tensor, stream_handle: int) -> tuple: + def fused_submit_args( + self, routes: torch.Tensor, stream_handle: int, *, spin_cycles: int + ) -> tuple: """Scheduler arguments of the bound GPU-direct copy's fused launch. The input and stream contract is ``submit()``'s. The copy endpoint launches this scheduler together with its copy, so no plan publication becomes pending and the launch triggers no programmatic dependency. + ``spin_cycles`` bounds this launch's cross-rank waits. """ + _check_spin_cycles(spin_cycles) with self._plan_channel_lock: if not self._gpu_direct_bound: raise RuntimeError("a fused launch requires a bound GPU-direct copy") @@ -583,17 +621,20 @@ def fused_submit_args(self, routes: torch.Tensor, stream_handle: int) -> tuple: "and reusable barrier epochs are shared" ) valid_tokens = self._check_routes(routes, None) - return self._halo_q_args(routes, valid_tokens * self.cfg.topk, enable_pdl=False) + return self._halo_q_args( + routes, valid_tokens * self.cfg.topk, enable_pdl=False, spin_cycles=spin_cycles + ) def _submit_routes( self, routes: torch.Tensor, stream_handle: int, *, - valid_tokens: int | None = None, + valid_tokens: int | None, spin_cycles: int, ) -> None: valid_tokens = self._check_routes(routes, valid_tokens) self._launch_op( routes, stream_handle, valid_route_count=valid_tokens * self.cfg.topk, + spin_cycles=spin_cycles, ) def _check_routes(self, routes: torch.Tensor, valid_tokens: int | None) -> int: @@ -618,6 +659,7 @@ def run( stream: object | None = None, *, validate_values: bool = False, + spin_cycles: int = SPIN_CYCLES, ) -> CudaSchedulerOutputs: if stream is not None and not self._plan_channel_bound: with torch.cuda.device(self.device): @@ -628,7 +670,7 @@ def run( self.stage(logical_expert_ids, validate_values=validate_values) else: self.stage(logical_expert_ids, validate_values=validate_values) - return self.launch(stream) + return self.launch(stream, spin_cycles=spin_cycles) def check_status(self) -> None: """Raise if any generation faulted since the last `reset_status()`. @@ -663,9 +705,12 @@ def reset_status(self) -> None: __all__ = [ + "ASSUMED_CLOCK_HZ", + "SPIN_CYCLES", "CudaPhysicalSlotScheduler", "CudaSchedulerConfig", "CudaSchedulerOutputs", "recommend_cuda_scheduler_ctas", + "spin_cycles_for_timeout", "symmetric_buffer_ints", ] diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py index cc74d5a6b34f..4b1b5ea133e0 100644 --- a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py @@ -629,13 +629,20 @@ def _require_live_bank_reusable(self) -> None: ) def submit_fused( - self, outputs: object, routes: torch.Tensor, route_epoch: torch.Tensor, route_target: int + self, + outputs: object, + routes: torch.Tensor, + route_epoch: torch.Tensor, + route_target: int, + *, + spin_cycles: int, ) -> HierarchicalCopyTicket: """Enqueue the bound scheduler and this GPU-direct copy as one kernel. Each scheduler CTA adds one to ``route_epoch`` once its routes are written, which makes it reach ``route_target``; the copy and its READY - publication follow in the same CTAs. + publication follow in the same CTAs. ``spin_cycles`` bounds the + launch's cross-rank waits and its wait for the route epoch. """ with self._lock: @@ -647,7 +654,9 @@ def submit_fused( raise RuntimeError("outputs are not from the bound CUDA scheduler") self._require_live_bank_reusable() try: - scheduler_args = self._scheduler.fused_submit_args(routes, self._stream) + scheduler_args = self._scheduler.fused_submit_args( + routes, self._stream, spin_cycles=spin_cycles + ) state = enter_cuda_stream(self._stream_owner) try: torch.ops.trtllm.moe_rebalance_tma_submit_fused( diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index a3b11a06bae1..f68780ac86bd 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -10,10 +10,11 @@ import importlib.util import math import os +import re import sys from contextlib import nullcontext from pathlib import Path -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace from typing import Any, List, Optional, Tuple import pytest @@ -23,6 +24,11 @@ _ROOT = Path(__file__).resolve().parents[4] _TORCH_ROOT = _ROOT / "tensorrt_llm/_torch" _QUEUE = "CUDA_SCALE_LAUNCH_QUEUES" +_SCHEDULER_RUNTIME = _TORCH_ROOT / "moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py" +# The scheduler's native bound and the bound of the default 1800 s MoE +# communication warmup timeout at the nominal 2 GHz clock. +_SERVING_SPIN_CYCLES = 8_000_000_000 +_WARMUP_SPIN_CYCLES = 1800 * 2_000_000_000 def _extract_class( @@ -252,8 +258,8 @@ def submit_fused(ctx, flags_mc, route_epoch, route_target, *scheduler_args): outputs = object() requests = [] - def fused_submit_args(routes, stream_handle): - requests.append((routes, stream_handle)) + def fused_submit_args(routes, stream_handle, *, spin_cycles): + requests.append((routes, stream_handle, spin_cycles)) return ("scheduler", routes) broadcast = broadcast_cls.__new__(broadcast_cls) @@ -273,9 +279,9 @@ def fused_submit_args(routes, stream_handle): broadcast.flags_mc = 4096 broadcast.terminal_flags = "terminals" - ticket = broadcast.submit_fused(outputs, "routes", "epoch", 8) + ticket = broadcast.submit_fused(outputs, "routes", "epoch", 8, spin_cycles=12345) - assert requests == [("routes", 19)] + assert requests == [("routes", 19, 12345)] assert launches == [("ctx", 4096, "epoch", 8, ("scheduler", "routes"))] # The launch runs on the copy stream and restores the caller's stream afterwards. owner = broadcast._stream_owner @@ -284,9 +290,9 @@ def fused_submit_args(routes, stream_handle): assert (broadcast._active_generation, broadcast._generation) == (1, 2) # The bank stays leased until the consumer of generation 1 is proven done. with pytest.raises(RuntimeError, match="still leased"): - broadcast.submit_fused(outputs, "routes", "epoch", 16) + broadcast.submit_fused(outputs, "routes", "epoch", 16, spin_cycles=12345) with pytest.raises(RuntimeError, match="bound CUDA scheduler"): - broadcast.submit_fused(object(), "routes", "epoch", 16) + broadcast.submit_fused(object(), "routes", "epoch", 16, spin_cycles=12345) @pytest.mark.parametrize( @@ -1219,6 +1225,9 @@ def test_fused_submit_orders_generations_and_defers_the_route_wait(auxiliary_sms route_epoch = object() group._route_epoch = route_epoch group._route_epoch_ptr = 77 + # Generation 1 launches during warmup, generation 2 while serving. + bounds = iter((_WARMUP_SPIN_CYCLES, _SERVING_SPIN_CYCLES)) + group._spin_bound = SimpleNamespace(cycles=lambda: next(bounds)) def driver_call(*entry): trace.append(entry) @@ -1235,10 +1244,13 @@ def driver_call(*entry): ) outputs = SimpleNamespace(physical_slot_ids=_Routes(8192, trace)) - def submit_fused(result, routes, epoch, target): + launch_bounds = [] + + def submit_fused(result, routes, epoch, target, *, spin_cycles): assert result is outputs and epoch is route_epoch generation = group._scheduled_generation trace.append(("FUSED", routes.shape[0], target, generation)) + launch_bounds.append(spin_cycles) return SimpleNamespace(generation=generation) group.scheduler = SimpleNamespace(device="cuda:0", outputs=outputs) @@ -1273,6 +1285,7 @@ def submit_fused(result, routes, epoch, target): assert trace.index(("lease", generation - 1)) < trace.index( ("FUSED", rows, target, generation) ) + assert launch_bounds == [_WARMUP_SPIN_CYCLES, _SERVING_SPIN_CYCLES] # The bound-owner fast path still rejects another thread, stream or device. state.thread = 8 @@ -1300,6 +1313,204 @@ def submit_fused(result, routes, epoch, target): assert group._plan_part is None +def _module_constants(path: Path, names: set[str]) -> dict: + tree = ast.parse(path.read_text()) + constants = { + node.targets[0].id: ast.literal_eval(node.value) + for node in tree.body + if isinstance(node, ast.Assign) + and len(node.targets) == 1 + and isinstance(node.targets[0], ast.Name) + and node.targets[0].id in names + } + assert set(constants) == names + return constants + + +def _scheduler_spin_namespace() -> dict: + namespace = _module_constants(_SCHEDULER_RUNTIME, {"SPIN_CYCLES", "ASSUMED_CLOCK_HZ"}) + _extract_function(_SCHEDULER_RUNTIME, "spin_cycles_for_timeout", namespace) + _extract_function(_SCHEDULER_RUNTIME, "_check_spin_cycles", namespace) + return namespace + + +def _moe_comm_timeout_guard_module(monkeypatch) -> ModuleType: + """Execute the timeout guard module without its torch and logger imports.""" + path = _TORCH_ROOT / "moe/fused_moe/moe_comm_timeout_guard.py" + tree = ast.parse(path.read_text()) + tree.body = [ + node + for node in tree.body + if not ( + (isinstance(node, ast.Import) and node.names[0].name == "torch") + or (isinstance(node, ast.ImportFrom) and node.module == "tensorrt_llm.logger") + ) + ] + module = ModuleType("moe_comm_timeout_guard_under_test") + module.logger = SimpleNamespace(info=lambda *args, **kwargs: None) + module.torch = None + # Dataclasses resolve their module through sys.modules. + monkeypatch.setitem(sys.modules, module.__name__, module) + exec(compile(tree, str(path), "exec"), module.__dict__) + return module + + +def test_scheduler_spin_cycles_convert_timeouts_at_the_nominal_clock(): + namespace = _scheduler_spin_namespace() + to_cycles = namespace["spin_cycles_for_timeout"] + + assert namespace["SPIN_CYCLES"] == _SERVING_SPIN_CYCLES + assert to_cycles(None) == _SERVING_SPIN_CYCLES + assert to_cycles(1800) == _WARMUP_SPIN_CYCLES + for seconds in (0, -1, 2.5, True): + with pytest.raises(ValueError, match="positive exact int"): + to_cycles(seconds) + for cycles in (0, -1, 1.0, True): + with pytest.raises(ValueError, match="positive exact int"): + namespace["_check_spin_cycles"](cycles) + + +@pytest.mark.parametrize( + ("environ", "warmup_cycles", "serving_cycles"), + [ + pytest.param({}, _WARMUP_SPIN_CYCLES, _SERVING_SPIN_CYCLES, id="defaults"), + pytest.param( + { + "TRTLLM_MOE_COMM_WARMUP_TIMEOUT_SEC": "600", + "TRTLLM_MOE_COMM_TIMEOUT_SEC": "30", + }, + 600 * 2_000_000_000, + 30 * 2_000_000_000, + id="configured", + ), + ], +) +def test_scheduler_spin_bound_follows_warmup_transitions_and_autotuning( + monkeypatch, environ, warmup_cycles, serving_cycles +): + guard_module = _moe_comm_timeout_guard_module(monkeypatch) + guard = guard_module.MoECommTimeoutGuard(environ=environ, is_capturing=lambda: False) + runtime = SimpleNamespace(**_scheduler_spin_namespace()) + monkeypatch.setitem( + sys.modules, + "tensorrt_llm._torch.moe.kernels.megamoe_scheduler.cuda_scheduler.runtime", + runtime, + ) + bound_cls = _extract_class( + _TORCH_ROOT / "moe/fused_moe/mega_moe/rebalance_slot_scheduler.py", + "SchedulerSpinBound", + None, + {"Callable": Any, "Optional": Optional}, + ) + tuning = [False] + bound = bound_cls(warmup_seconds=guard.budgets().warmup_seconds, is_tuning=lambda: tuning[0]) + guard.register(bound) + + # The engine starts in serving; its warmup transitions switch the bound. + assert bound.cycles() == serving_cycles + guard.set_warmup(True) + assert bound.cycles() == warmup_cycles + guard.set_warmup(False) + assert bound.cycles() == serving_cycles + # AutoTuner profiling outside warmup gets the warmup bound too. + tuning[0] = True + assert bound.cycles() == warmup_cycles + # The startup warmup launch always uses the warmup bound. + assert bound.warmup_cycles == warmup_cycles + + +def test_scheduler_launches_pass_their_spin_bound_to_the_kernel(): + # Locate spin_cycles in the registered operator schema rather than hard-coding it. + op_source = ( + _ROOT / "cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.cpp" + ).read_text() + schema = "".join(re.findall(r'"([^"]*)"', op_source.split("m.def(", 1)[1].split(");", 1)[0])) + parameters = re.fullmatch(r"moe_rebalance_halo_q\((.*)\) -> \(\)", schema).group(1) + names = [parameter.split()[-1] for parameter in parameters.split(",")] + spin_index = names.index("spin_cycles") + + launched = [] + torch = SimpleNamespace( + Tensor=_Routes, + int32="int32", + ops=SimpleNamespace( + trtllm=SimpleNamespace(moe_rebalance_halo_q=lambda *args: launched.append(args)) + ), + ) + namespace = _scheduler_spin_namespace() + namespace.update( + { + "torch": torch, + "enter_cuda_stream": lambda stream: stream, + "exit_cuda_stream": lambda state: None, + "_stream_handle": lambda stream: int(getattr(stream, "cuda_stream", stream)), + } + ) + scheduler_cls = _extract_class( + _SCHEDULER_RUNTIME, + "CudaPhysicalSlotScheduler", + {"launch", "_launch_op", "_halo_q_args", "fused_submit_args", "_check_routes"}, + namespace, + ) + scheduler = scheduler_cls.__new__(scheduler_cls) + scheduler.cfg = SimpleNamespace( + ep_size=8, + logical_expert_count=384, + extra_slots_per_rank=4, + local_rank=0, + route_count=8192 * 6, + ctas=8, + threads=512, + enable_pdl=True, + topk=6, + max_tokens_per_rank=8192, + ) + scheduler.outputs = SimpleNamespace( + physical_slot_ids="slots", + hot_expert_ids="ids", + hot_expert_group_level="levels", + hot_expert_source_ranks="owners", + ) + for name in ( + "peer_base", + "status", + "partial", + "route_aux", + "grid_sync", + "plan_workspace", + "route_prefix", + ): + setattr(scheduler, name, name) + scheduler.logical_expert_ids = "staged routes" + scheduler.device = "cuda:0" + scheduler._connected = True + scheduler._plan_channel_lock = nullcontext() + scheduler._plan_channel_bound = False + scheduler._plan_channel_pending = False + scheduler._plan_channel_ptr = 0 + scheduler._plan_channel_abi_version = 0 + scheduler._plan_channel_words = 0 + scheduler._plan_channel_route_features = 0 + scheduler._plan_channel_torch_stream = SimpleNamespace(cuda_stream=19) + scheduler._launch_stream_handle = None + + # The standalone kernel keeps the native bound unless the caller passes one. + scheduler.launch(19) + scheduler.launch(19, spin_cycles=_WARMUP_SPIN_CYCLES) + assert [args[spin_index] for args in launched] == [_SERVING_SPIN_CYCLES, _WARMUP_SPIN_CYCLES] + + # The fused launch takes the same scheduler arguments with the caller's bound. + scheduler._gpu_direct_bound = True + scheduler._plan_channel_stream_handle = 19 + routes = _Routes(16, []) + routes.is_cuda = True + args = scheduler.fused_submit_args(routes, 19, spin_cycles=_WARMUP_SPIN_CYCLES) + assert args[spin_index] == _WARMUP_SPIN_CYCLES + assert args[spin_index - 1] is False # No programmatic dependency in the fused kernel. + with pytest.raises(ValueError, match="positive exact int"): + scheduler.fused_submit_args(routes, 19, spin_cycles=0) + + def test_raw_stream_switch_restores_the_previous_stream_and_device(): calls = [] current = {"device": 0, "streams": {0: (10, 0, 1), 3: (30, 3, 1)}} diff --git a/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py b/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py index a6a6b576d443..a8c8fc6b698f 100644 --- a/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py +++ b/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py @@ -461,6 +461,27 @@ def test_helper_slot_mode_rejects_unknown_mode(): PerLayerDynamicMoeLoadBalancer(ep_rank=0, ep_size=4, config=config) +def test_scheduler_spin_bound_follows_the_warmup_guard_until_shutdown(monkeypatch): + from tensorrt_llm._torch.moe.fused_moe import moe_comm_timeout_guard + + guard = moe_comm_timeout_guard.MoECommTimeoutGuard(environ={}, is_capturing=lambda: False) + monkeypatch.setattr(moe_comm_timeout_guard, "_DEFAULT_GUARD", guard) + manager = _manager() + + bound = manager.get_scheduler_spin_bound() + + assert manager.get_scheduler_spin_bound() is bound + assert bound.cycles() == 8_000_000_000 + guard.set_warmup(True) + assert bound.cycles() == 1800 * 2_000_000_000 + guard.set_warmup(False) + assert bound.cycles() == 8_000_000_000 + assert manager.shutdown() + guard.set_warmup(True) + # A released bound no longer follows the guard. + assert bound.cycles() == 8_000_000_000 + + def test_per_layer_dynamic_manager_shares_the_base_lifecycle_only(): manager = _manager() assert isinstance(manager, BaseMoeLoadBalancer) From b6b845d9d854b9d01886bd9693da926388ccaefc Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 07:59:17 -0700 Subject: [PATCH 19/59] [None][chore] Align a copyright header and remove history narration from comments megamoe_shared_mlp.py was the only file in tensorrt_llm/_torch/modules/ with the pre-SPDX header line; switch it to the SPDX form that its siblings use. Rewrite comments in the HALO-Q kernel, the TMA copy header, the MegaMoE CuTe DSL custom op and a MoE backend test that compared the code with earlier versions of itself ("the old linear scan", "used to run", "as before"), so that they describe the code as it stands. Comment-only change. Signed-off-by: jintaop --- .../dynamicEplb/moeRebalanceHaloQKernel.cuh | 53 +++++++++---------- .../loadBalance/dynamicEplb/moeRebalanceTma.h | 2 +- .../_torch/modules/megamoe_shared_mlp.py | 5 +- .../custom_ops/cute_dsl_megamoe_custom_op.py | 6 +-- tests/unittest/_torch/moe/test_moe_backend.py | 4 +- 5 files changed, 33 insertions(+), 37 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh index 9f4733972f65..c1053da8e8d2 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh @@ -649,12 +649,12 @@ __device__ __forceinline__ int plan_halo_m_register8_impl( continue; } - /* owner and owner_next are now warp-uniform, so every lane can read + /* owner and owner_next are warp-uniform here, so every lane can read * candidates[] and histogram[] directly. Both are written by * prepare_halo_m_sparse8 before its bar.sync and never inside this - * loop, so a uniform load returns exactly what the owner would have - * broadcast. The increment stays owner-only and stays ahead of the - * mass test, as before. */ + * loop, so every lane loads the same values and no broadcast from the + * owner is needed. Only the owner advances lane_next, and it does so + * before the mass test. */ int const expert = candidates[owner * plan_helpers + owner_next]; if (lane == owner) ++lane_next; @@ -692,11 +692,10 @@ __device__ __forceinline__ int plan_halo_m_register8_impl( if constexpr (FixedEp == 8 && FixedHelpers == 4) { int slack = lane != owner ? imax(0, target - lane_load) : 0; - /* Two independent full-mask sums instead of one lane-varying-mask - * reduce followed by two broadcasts. Integer addition is exact, so - * masking the addend gives the same half sums; the two reduces do - * not depend on each other, so the chain is depth 1 instead of 3, - * and the redux.sync no longer runs with a per-lane membermask. */ + /* Two independent full-mask sums give the two half sums. Integer + * addition is exact, so zeroing the other half's addends yields + * each half sum exactly; the two reduces do not depend on each + * other, and neither redux.sync uses a per-lane membermask. */ int room0 = __reduce_add_sync(mask, lane < 4 ? slack : 0); int room4 = __reduce_add_sync(mask, lane >= 4 ? slack : 0); unsigned receive_mask = __ballot_sync(mask, lane_recv < plan_helpers); @@ -868,7 +867,7 @@ __device__ __forceinline__ int plan_halo_m_register8( } /* Exact HALO-M state machine executed by warp 0. Each lane owns one rank; - * warp reductions replace all serial rank/group/destination scans. */ + * rank/group/destination scans are warp reductions rather than serial loops. */ __device__ int plan_halo_m_warp(int const* histogram, int* load, int* quota, int const* candidates, int* recv_used, int* send_used, int* next_index, int* frozen, int* plan, int ep, int experts, int helpers) { @@ -1233,7 +1232,7 @@ __device__ __forceinline__ void halo_q_direction8( * * `b`, `hosts` and `expert` are the lane's own broadcast identity and do not * change between phases or rounds, so the caller hoists them. Reading them - * here put a four-deep dependent load chain (order[lane] -> b -> host_mask[b] + * here would put a four-deep dependent load chain (order[lane] -> b -> host_mask[b] * and the GLOBAL p_expert[b] -> expert -> counts[...]) at the head of all * fourteen phases. */ __device__ __forceinline__ void halo_q_pair8( @@ -1281,13 +1280,9 @@ __device__ __forceinline__ void scheduler_tail_barrier() /* Build order[] and host_mask[] with one lane per broadcast. * - * This replaces serial insertion-sort and mask-construction loops while - * preserving the same stable order. - * - * order[]: lane b's position is the number of broadcasts that sort before it. - * The `j < lane` tiebreak makes this the same stable order the - * insertion sort produced (its shift test is a strict `>`), so the - * result is bit-identical, not merely equivalent. + * order[]: lane b's position is the number of broadcasts that sort before it, + * ordered by expert. The `j < lane` tiebreak orders equal experts by + * broadcast index, so the order is stable. * host_mask: the destination group is a contiguous rank range, so the mask is * closed form; no loop over group members is needed. * rank_group_mask handles the full EP32 word without a shift by32. @@ -1354,8 +1349,9 @@ __device__ void run_halo_q(int const* counts, int* quota, int* plan, int broadca if (broadcasts >= 2) { int const ring_size = plan_ep - 1; - /* Use two unconditional repair rounds. Phases with no required repair - * are no-ops, so this preserves the conditional algorithm's result. */ + /* Run two repair rounds unconditionally. A phase with no required + * repair is a no-op, so the plan equals running only the required + * repairs. */ int const repair_rounds = 2; for (int round = 0; round < repair_rounds; ++round) { @@ -1394,7 +1390,6 @@ __device__ void run_halo_q(int const* counts, int* quota, int* plan, int broadca } } } - /* The unconditional repair rounds above replace the conditional scan. */ return; } if (warp == 0) @@ -1515,12 +1510,12 @@ __device__ void build_route_prefix(int const* counts, int const* quota, int* pre int source_count = counts[local_rank * experts + expert]; int source_diagonal = imin(source_count, quota[b * ep + local_rank]); int residual_supply = source_count - source_diagonal; - /* The loop used to run to destination == ep. route_destination only - * ever loads prefix[destination + step] under `candidate < ep`, so the - * ep entry is never read, and the accumulation at ep-1 fed nothing but - * that dead entry. Both are gone; entry 0 is still written because it - * costs one store and leaving it stale would depend on nobody ever - * reading it, which is a weaker guarantee than writing the right value. */ + /* route_destination loads prefix[destination + step] only under + * `candidate < ep`, so entry ep is never read: the loop stops before + * destination == ep and skips the accumulation after the last + * destination, which would only feed that entry. Entry 0 is written + * although route_destination never reads it: it costs one store and + * leaves no stale entry below ep. */ int destination_before = 0; for (int destination = 0; destination < ep; ++destination) { @@ -1591,8 +1586,8 @@ __device__ bool wait_for_plan_epoch(int* marker, Epoch epoch, unsigned long long return *ready != 0; } -/* Prefix entries are monotonic. Binary lifting returns the same number of - * thresholds <= ordinal as the old linear scan, including repeated entries. */ +/* Prefix entries are monotonic, so binary lifting returns the number of + * thresholds <= ordinal, counting repeated entries. */ __device__ __forceinline__ int route_destination(int const* prefix, int ordinal, int ep) { int destination = 0; diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h index dffa0e523298..582e6b3c6594 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h @@ -122,7 +122,7 @@ extern "C" * Slice j spans [j*8192,min((j+1)*8192,bytes)) in its virtual byte stream and can * cross any number of segments. It keeps ONE GPU/CTA/warp owner throughout. * This rank owns floor((first_slice+j*slice_stride)/slice_divisor). Uniform - * routing uses divisor=1, reciprocal=0 and retains the original integer path. + * routing uses divisor=1, reciprocal=0 and skips the division. * Otherwise reciprocal=UINT64_MAX/divisor enables an exact mulhi+correction. * prefix_end compacts rank-local slices across ranges for SM/worker round-robin. * No reset at weight boundaries. diff --git a/tensorrt_llm/_torch/modules/megamoe_shared_mlp.py b/tensorrt_llm/_torch/modules/megamoe_shared_mlp.py index e041cd137e50..4aea1efc61d8 100644 --- a/tensorrt_llm/_torch/modules/megamoe_shared_mlp.py +++ b/tensorrt_llm/_torch/modules/megamoe_shared_mlp.py @@ -1,10 +1,11 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, diff --git a/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py b/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py index de498b2cd01e..ed32d68363b6 100644 --- a/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py +++ b/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py @@ -1228,7 +1228,7 @@ def _expand_megamoe_ready_tactics( choices = _token_back_ready_granularity_choices(load_balance=load_balance) if sm_version == 107 and choices != ("expert",): # At every token bucket include a legal local-TMA/atomic-counter - # base; the historical candidates otherwise exclude token_tile. + # base; the curated SM107 candidates otherwise exclude token_tile. for cluster, fallback, policy, mode, flag, epi in ( ([2, 1, 1], None, ("grouped", 512), "reuse_dispatch_warps", 4, (1, 1)), ([2, 1, 1], None, ("grouped", 512), "reuse_dispatch_warps", 4, (2, 4)), @@ -3004,7 +3004,7 @@ def run_cute_dsl_megamoe_nvfp4_blackwell( # ``hot_expert_weight_ready_flags`` is the local uint64[EP] terminal # view published by the weight transport, ``..._generation`` the # expected generation g in [1, 2**63). All three default to the - # unchanged S == 0 baseline. + # S == 0 kernel without helper slots. helper_expert_count: int = 0, hot_expert_weight_ready_flags: Optional[torch.Tensor] = None, hot_expert_weight_ready_generation: int = 0, @@ -3207,7 +3207,7 @@ def _( # ``hot_expert_weight_ready_flags`` is the local uint64[EP] terminal # view published by the weight transport, ``..._generation`` the # expected generation g in [1, 2**63). All three default to the - # unchanged S == 0 baseline. + # S == 0 kernel without helper slots. helper_expert_count: int = 0, hot_expert_weight_ready_flags: Optional[torch.Tensor] = None, hot_expert_weight_ready_generation: int = 0, diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index 91e1b359a738..27486ffec9c2 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -2173,8 +2173,8 @@ def test_megamoe_compute_slots_do_not_widen_resident_expert_axis(): assert moe.compute_slot_count == 11 assert moe.expert_size_per_partition == 8 - # Catch the old representation (expert_size_per_partition == H + S) - # instead of silently adding S a second time. + # A resident axis that already counts the helper slots + # (expert_size_per_partition == H + S) must raise instead of adding S again. moe.expert_size_per_partition = 11 with pytest.raises(RuntimeError, match="resident expert geometry changed"): _ = moe.compute_slot_count From b39f848d72bd258eca558b87696622a0fdd1d2da Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 08:14:36 -0700 Subject: [PATCH 20/59] [None][chore] Describe the dynamic EPLB kernels as they stand and drop a duplicate CPU test entry Reword three comments in the dynamic EPLB kernels that compared the code with earlier versions. Remove the explicit l0_cpu entry for test_megamoe_shared_fc12_autotune.py: the unittest/_torch/moe directory entry already collects it, so the file ran twice on the CPU stage. Signed-off-by: jintaop --- .../moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh | 4 ++-- .../moe/loadBalance/dynamicEplb/moeRebalancePlanPolicy.h | 2 +- .../kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu | 2 +- tests/integration/test_lists/test-db/l0_cpu.yml | 1 - 4 files changed, 4 insertions(+), 5 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh index c1053da8e8d2..970fac3420d4 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh @@ -256,8 +256,8 @@ __device__ void local_histogram_pass(int const* routes, int* partial, int route_ int const begin = cta_begin + warp * warp_span; int const end = imin(cta_end, begin + warp_span); /* Issue a small load group before consuming it to expose memory-level - * parallelism. The binning is unchanged, and atomicAdd regrouping is - * safe because integer addition commutes. */ + * parallelism. Batching does not change which bin a route lands in, and + * atomicAdd order does not matter because integer addition commutes. */ constexpr int kBatch = 4; for (int base_index = begin; base_index < end; base_index += 32 * kBatch) { diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlanPolicy.h b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlanPolicy.h index b577ad28d2cc..ccb2f820b867 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlanPolicy.h +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlanPolicy.h @@ -202,7 +202,7 @@ __device__ void publish_plan_channel(int* channel, int const* ids, int const* le /* Small ABI6 plans fit in one warp. Load each record once and compare its * register values through shuffles. Two mode bits per helper fit in a register - * for S<=16. Larger/malformed plans use the unchanged scalar publisher. + * for S<=16. Larger or malformed plans use the scalar publisher. * Only lane 0 writes mapped words; its SYS release orders every field. */ template __device__ void publish_plan_channel_warp(int* channel, int const* ids, int const* levels, int const* owners, diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu index a2047a3b9ca5..d2d5fe82c6d6 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu @@ -484,7 +484,7 @@ __device__ __forceinline__ void consume_bank(MegamoeTmaCopySegment const* segmen } } // These exact scalar intervals are disjoint from ALL TMA writes. - // No intermediate group drain; final full-write wait/fence is unchanged. + // No intermediate group drain: the final full-write wait and fence cover them. scalar_words(f, 0, f.scalar_prefix); scalar_words(f, f.scalar_prefix + f.bulk_bytes, f.bytes); } diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index eda6293ff984..a02f527b12d2 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -48,7 +48,6 @@ l0_cpu: - unittest/_torch/kernels/kimi_k3/test_k3_mla_decode_view.py - unittest/_torch/kernels/kimi_k3/test_k3_mla_launch.py - unittest/_torch/kernels/kimi_k3/test_k3_kda_pool_layout.py - - unittest/_torch/moe/test_megamoe_shared_fc12_autotune.py - unittest/_torch/compilation/test_auto_multi_stream.py - unittest/_torch/compilation/test_phase_selective_forward.py - unittest/_torch/compilation/test_remove_copy_pass.py::test_remove_copy_preserves_minimax_producer_outputs From 14fabccec6629df0c6858a67470247a2981ea341 Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 14:02:53 -0700 Subject: [PATCH 21/59] [None][test] Run the HALO-Q spin-bound test without a GPU The spin bound asks AutoTuner whether kernels are being profiled, and creating the AutoTuner singleton reads the CUDA device name, which fails on the CPU-only l0_cpu stage. Stub AutoTuner.get in the test, and check that a profiling AutoTuner gets the warmup bound. Signed-off-by: jintaop --- tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py b/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py index a8c8fc6b698f..065f332911b6 100644 --- a/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py +++ b/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py @@ -462,16 +462,23 @@ def test_helper_slot_mode_rejects_unknown_mode(): def test_scheduler_spin_bound_follows_the_warmup_guard_until_shutdown(monkeypatch): + from tensorrt_llm._torch.autotuner import AutoTuner from tensorrt_llm._torch.moe.fused_moe import moe_comm_timeout_guard guard = moe_comm_timeout_guard.MoECommTimeoutGuard(environ={}, is_capturing=lambda: False) monkeypatch.setattr(moe_comm_timeout_guard, "_DEFAULT_GUARD", guard) + # The real AutoTuner reads the CUDA device name; this test runs without a GPU. + tuner = SimpleNamespace(is_tuning_mode=False) + monkeypatch.setattr(AutoTuner, "get", lambda: tuner) manager = _manager() bound = manager.get_scheduler_spin_bound() assert manager.get_scheduler_spin_bound() is bound assert bound.cycles() == 8_000_000_000 + tuner.is_tuning_mode = True + assert bound.cycles() == 1800 * 2_000_000_000 + tuner.is_tuning_mode = False guard.set_warmup(True) assert bound.cycles() == 1800 * 2_000_000_000 guard.set_warmup(False) From f98b30cd19ada1c5170c28b1ffd99e83d5567cc8 Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 14:43:28 -0700 Subject: [PATCH 22/59] [None][perf] Keep dynamic EPLB routes alive without an allocator stream event plan_schedule() protected the logical routes that HALO-Q reads on the copy stream with record_stream(). When such a block is freed, the caching allocator records an event on the copy stream and then queries it on every later allocation until the copy stream reaches it. The copy stream runs behind the host by up to an iteration, so on GB200 rank 0 every iteration paid about 3,400 cudaEventQuery calls (6 ms of API time under nsys), and the OFF arm none. The scheduler group now holds the routes until wait_for_routes() queues MAIN's wait on the route epoch. Each HALO-Q CTA reads its routes before it raises that epoch with a system-scope release, so once MAIN has queued the wait, a block the allocator returns to MAIN can only be reused by work queued after it. Per-layer dynamic EPLB already rejects CUDA graphs, and the routes are produced on MAIN. wait_for_routes() also returns at once for a generation whose wait is already queued; MegaMoE calls it up to three times per layer. Signed-off-by: jintaop --- .../mega_moe/rebalance_slot_scheduler.py | 35 ++++++++++++++----- .../_torch/moe/test_dynamic_eplb_contracts.py | 10 ++++++ 2 files changed, 36 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py index dac4279e21f9..7d1d19d90e83 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py @@ -165,6 +165,8 @@ def __init__( self._finished_generation = 0 self._scheduled_generation = 0 self._route_wait_generation = 0 + # Routes of the latest launch, held until MAIN waits for its route epoch. + self._routes_in_flight: Optional[torch.Tensor] = None self._owner_thread_id: Optional[int] = None self._execution_stream_handle: Optional[int] = None self._closed = False @@ -293,6 +295,7 @@ def close(self) -> None: self._execution_stream_handle = None self._plan_part = None self._pending_release = None + self._routes_in_flight = None self.lease = None self.broadcaster = None self.scheduler = None @@ -390,9 +393,14 @@ def _wait_route_epoch(self, stream_handle: int) -> None: def plan_schedule(self, logical_expert_ids: torch.Tensor) -> Tuple[Any, int]: """Submit contiguous CUDA int32 routes [tokens, topk] without staging. - The route producer must run on MAIN and must not overwrite this tensor - until HALO-Q finishes. record_stream protects allocator reuse while COPY - consumes the input. Empty ranks still submit and join the EP exchange. + The route producer must allocate and write this tensor on MAIN and must + not overwrite it until HALO-Q finishes. The group holds the tensor until + wait_for_routes() orders MAIN after the route epoch. Every HALO-Q CTA + reads its routes before it raises that epoch, so the caching allocator + can only reuse the block for MAIN work queued after that wait. This needs + no COPY-stream allocator event, which every later allocation would poll + while COPY runs behind MAIN. Empty ranks still submit and join the EP + exchange. """ self._check_owner() if self._plan_part is not None or self._generation != self._finished_generation: @@ -412,7 +420,7 @@ def plan_schedule(self, logical_expert_ids: torch.Tensor) -> Tuple[Any, int]: # Both waits capture their event's current record before a later plan # can re-record it. No host synchronization or temporary event is needed. - logical_expert_ids.record_stream(self.copy_stream) + self._routes_in_flight = logical_expert_ids self._record_event(self._input_ready_handle, self._execution_stream_handle) self._wait_event(self._stream_handle, self._input_ready_handle) outputs = self.scheduler.outputs @@ -452,13 +460,22 @@ def plan_finish( return physical, self._generation def wait_for_routes(self) -> None: - """Order the first MAIN route read after HALO-Q, without a host wait.""" + """Order the first MAIN route read after HALO-Q, without a host wait. + + The first call of a generation checks the owner and queues the stream + wait. Later calls of that generation return at once: their consumers + are already ordered after it on MAIN. + """ + generation = self._generation + if generation > 0 and generation == self._route_wait_generation and not self._closed: + return self._check_owner() - if self._generation <= 0: + if generation <= 0: raise RuntimeError("No scheduler generation is available") - if self._route_wait_generation != self._generation: - self._wait_route_epoch(self._execution_stream_handle) - self._route_wait_generation = self._generation + self._wait_route_epoch(self._execution_stream_handle) + self._route_wait_generation = generation + # MAIN work queued from here on runs after every HALO-Q read of the routes. + self._routes_in_flight = None def plan( self, logical_expert_ids: torch.Tensor, *, defer_wait: bool = False diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index f68780ac86bd..5f076ae7490c 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -1273,10 +1273,15 @@ def submit_fused(result, routes, epoch, target, *, spin_cycles): assert physical.shape == (rows, 6) assert actual_generation == generation assert route_wait not in trace + # The group, not an allocator stream event, keeps the routes alive + # until MAIN waits for every HALO-Q read of them. + assert group._routes_in_flight is routes group.wait_for_routes() + assert group._routes_in_flight is None group.wait_for_routes() group.finish() assert trace.count(route_wait) == 1 + assert not any(entry[0] == "lifetime" for entry in trace) assert trace.index(("wait", 19, 1)) < trace.index(("FUSED", rows, target, generation)) assert trace.index(("FUSED", rows, target, generation)) < trace.index(route_wait) assert trace[-2:] == [("record", 3, 44), ("consumer_release", generation, 3)] @@ -1312,6 +1317,11 @@ def submit_fused(result, routes, epoch, target, *, spin_cycles): assert not any(entry[0] == "FUSED" for entry in trace) assert group._plan_part is None + # A closed group rejects every route wait, including one of a waited generation. + group._closed = True + with pytest.raises(RuntimeError, match="closed"): + group.wait_for_routes() + def _module_constants(path: Path, names: set[str]) -> dict: tree = ast.parse(path.read_text()) From 22c30067ef0fd68f3660a15f2e292567254f0f7d Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 14:45:22 -0700 Subject: [PATCH 23/59] [None][perf] Check the bound EPLB reuse authority by identity per layer The copy endpoint validates its generation reuse authority when the authority is bound, and binding is single-assignment. Each layer's submit re-ran that validation twice, once to prove the previous generation reusable and once before the fused launch. Compare the bound authority by identity on that path instead. Also import MoEReplicaPlan once at module scope rather than on every route call. Signed-off-by: jintaop --- .../moe/fused_moe/per_layer_dynamic_eplb.py | 3 +-- .../megamoe_scheduler/sami/hierarchical.py | 5 ++-- .../_torch/moe/test_dynamic_eplb_contracts.py | 24 ++++++++++++++++--- 3 files changed, 25 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py index f8fe733bc653..56604ce5f690 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py @@ -16,6 +16,7 @@ import torch +from .impl_contract import MoEReplicaPlan from .moe_load_balancer import BaseMoeLoadBalancer @@ -301,8 +302,6 @@ def route( generation = int(generation) self._arena.publish_ready_generation(generation) - from .impl_contract import MoEReplicaPlan - plan = MoEReplicaPlan( resident_experts_per_rank=self.resident_experts_per_rank, compute_slots_per_rank=self.compute_slots_per_rank, diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py index 4b1b5ea133e0..93c6cd0eee1c 100644 --- a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py @@ -605,9 +605,10 @@ def mark_collective_reuse_safe( with self._lock: if self._handoff_aborted or self._scheduler is None: raise RuntimeError("hierarchical scheduler handoff is not operational") + # bind_generation_reuse_authority() validated this single-assignment + # authority, so identity suffices on the per-layer path. if authority is not self._reuse_authority: raise PermissionError("generation reuse authority identity mismatch") - self._validate_reuse_authority(authority) if generation != self._active_generation: raise ValueError("generation is not the active live bank") if self._local_release_generation != generation: @@ -617,9 +618,9 @@ def mark_collective_reuse_safe( self._collective_safe_generation = generation def _require_live_bank_reusable(self) -> None: + # The bound authority was validated once, when it was bound. if self._reuse_authority is None: raise RuntimeError("single live bank has no bound collective reuse authority") - self._validate_reuse_authority(self._reuse_authority) if ( self._active_generation and self._collective_safe_generation < self._active_generation diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index 5f076ae7490c..b53b80cf6253 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -247,7 +247,7 @@ def submit_fused(ctx, flags_mc, route_epoch, route_target, *scheduler_args): broadcast_cls = _extract_class( source, "HierarchicalSamiWeightBroadcast", - {"submit_fused", "_require_live_bank_reusable"}, + {"submit_fused", "_require_live_bank_reusable", "mark_collective_reuse_safe"}, { "torch": torch, "HierarchicalCopyTicket": ticket_cls, @@ -267,8 +267,14 @@ def fused_submit_args(routes, stream_handle, *, spin_cycles): broadcast._handoff_aborted = False broadcast._outputs = outputs broadcast._scheduler = SimpleNamespace(fused_submit_args=fused_submit_args) - broadcast._reuse_authority = object() - broadcast._validate_reuse_authority = lambda authority: None + authority = object() + broadcast._reuse_authority = authority + + def revalidate(_authority): + raise AssertionError("the bound authority was validated when it was bound") + + # The per-layer path checks the bound authority by identity only. + broadcast._validate_reuse_authority = revalidate broadcast._active_generation = 0 broadcast._collective_safe_generation = 0 broadcast._generation = 1 @@ -294,6 +300,18 @@ def fused_submit_args(routes, stream_handle, *, spin_cycles): with pytest.raises(RuntimeError, match="bound CUDA scheduler"): broadcast.submit_fused(object(), "routes", "epoch", 16, spin_cycles=12345) + # Only the bound authority, after the local consumer release, frees the bank. + with pytest.raises(PermissionError, match="identity mismatch"): + broadcast.mark_collective_reuse_safe(object(), 1) + broadcast._local_release_generation = 0 + with pytest.raises(RuntimeError, match="local consumer release"): + broadcast.mark_collective_reuse_safe(authority, 1) + broadcast._local_release_generation = 1 + broadcast.mark_collective_reuse_safe(authority, 1) + assert broadcast._collective_safe_generation == 1 + ticket = broadcast.submit_fused(outputs, "routes", "epoch", 16, spin_cycles=12345) + assert ticket.generation == 2 + @pytest.mark.parametrize( ( From 6460d105581b644bad0e80d7f4386428727c8000 Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 14:52:24 -0700 Subject: [PATCH 24/59] [None][perf] Reuse MegaMoE launch buffers and their cute views across layers Every MegaMoE launch rebuilt state that is fixed for a bucket, once per layer: - MegaMoeSymmMemProvider.get_regions() re-sliced the five symmetric regions; the provider now builds them once per inference-mode state (a view made under inference mode cannot be updated in place outside it). - The backend rebuilt the per-bucket _MegaMoeBuffers bundle; it is now built once per (bucket, top_k, hidden, device, dtype, inference mode). - The runner converted the activation, activation SF, top-k weights, top-k indices and combine output with from_dlpack every launch. When the backend launches on its staged buffers it now hands the op its provider's cute-view cache: the symmetric regions are keyed by every field a view encodes, the combine output keeps one slot because its live rows change per iteration but not per layer, and the module's staging indices join its weight cache. Tuning launches still convert afresh. - Workspace and READY-flag pointers are cached by (address, alignment), the only fields a cute pointer encodes; the stream wrapper comes from the raw current stream; unique_id() is built once per runner. - SM100 routes were widened to int64 by a separate cast kernel before being copied into the int64 staging buffer. The staging copy now widens them itself, so the first route reader, and the route wait, move to that copy. All of these produce the same kernel arguments as before. Signed-off-by: jintaop --- .../custom_ops/cute_dsl_megamoe_custom_op.py | 129 +++++++++++++++--- .../fused_moe/mega_moe/mega_moe_cute_dsl.py | 60 ++++++-- .../_torch/moe/test_dynamic_eplb_contracts.py | 123 ++++++++++++++++- tests/unittest/_torch/moe/test_moe_backend.py | 86 ++++++++++++ 4 files changed, 359 insertions(+), 39 deletions(-) diff --git a/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py b/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py index ed32d68363b6..fc4ea0aed7be 100644 --- a/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py +++ b/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py @@ -1551,19 +1551,60 @@ def to_cute_ptr(t, assumed_align=16): return to_cute, to_cute_ptr, SymBufferHost, cute + def _cute_view_key(t: torch.Tensor, assumed_align: int) -> Tuple: + """Everything a cute view of ``t`` encodes.""" + return (t.data_ptr(), tuple(t.shape), t.stride(), t.dtype, t.device.index, assumed_align) + def _to_cute_cached(cache: dict, t: torch.Tensor, assumed_align: int = 16): """Return the cute view of ``t`` from ``cache``, converting it on a miss. The key holds everything the view encodes, so a hit equals a fresh conversion even if the address was reused. Only for tensors whose - layout is fixed (module weights): one entry per tensor. + layout is fixed (module weights, launch buffers): one entry per tensor. """ - key = (t.data_ptr(), tuple(t.shape), t.stride(), t.dtype, t.device.index, assumed_align) + key = _cute_view_key(t, assumed_align) cute_tensor = cache.get(key) if cute_tensor is None: cute_tensor = cache[key] = _cute_launch_helpers()[0](t, assumed_align) return cute_tensor + def _to_cute_slot(cache: dict, slot: str, t: torch.Tensor, assumed_align: int = 16): + """Like ``_to_cute_cached`` but keep only the latest view under ``slot``. + + For a buffer whose live row count changes between iterations but not + between the layers of one iteration. + """ + key = _cute_view_key(t, assumed_align) + entry = cache.get(slot) + if entry is None or entry[0] != key: + entry = cache[slot] = (key, _cute_launch_helpers()[0](t, assumed_align)) + return entry[1] + + # cute pointers by (address, alignment): a pointer encodes nothing else, so + # a hit equals a fresh ``make_ptr``. Bounded because tuning visits many + # workspaces; a pointer holds no tensor reference. + _MEGAMOE_CUTE_PTR_CACHE: dict = {} + _MEGAMOE_CUTE_PTR_CACHE_MAX_ENTRIES = 256 + + def _to_cute_ptr_cached(t: torch.Tensor, assumed_align: int = 16): + key = (t.data_ptr(), assumed_align) + ptr = _MEGAMOE_CUTE_PTR_CACHE.get(key) + if ptr is None: + if len(_MEGAMOE_CUTE_PTR_CACHE) >= _MEGAMOE_CUTE_PTR_CACHE_MAX_ENTRIES: + _MEGAMOE_CUTE_PTR_CACHE.clear() + ptr = _MEGAMOE_CUTE_PTR_CACHE[key] = _cute_launch_helpers()[1](t, assumed_align) + return ptr + + # cuda-python stream wrappers by raw handle. + _MEGAMOE_CU_STREAMS: dict = {} + + def _current_cu_stream(): + handle = torch._C._cuda_getCurrentRawStream(torch._C._cuda_getDevice()) + stream = _MEGAMOE_CU_STREAMS.get(handle) + if stream is None: + stream = _MEGAMOE_CU_STREAMS[handle] = cuda.CUstream(handle) + return stream + def _evict_latched_tuning_workspaces() -> None: """Pop every tuning-latched local workspace and prune its fence keys. @@ -1654,7 +1695,7 @@ class MegaMoeSymmRegions: # (max_T, 1, hidden): both reduction forms collapse the top-k axis in-op. combine_output: torch.Tensor shared_workspace: torch.Tensor # (shared_ws_bytes,) uint8 - peer_offsets: List[int] # symmetric peer-pointer deltas + peer_offsets: Tuple[int, ...] # symmetric peer-pointer deltas rank: int world_size: int @@ -1775,6 +1816,13 @@ def __init__( peer_ptr = int(self._handle.buffer_ptrs[r]) self.peer_offsets.append(peer_ptr - local_base) + # Built on first use per inference-mode state: a view made under + # inference mode cannot be updated in place outside it. + self._regions: dict = {} + # Cute views of these regions, shared by every layer launching on + # this provider (set as the op's buffer cache by the backend). + self.cute_views: dict = {} + logger.debug( "[MegaMoeSymmMemProvider] group=%s rank=%d/%d total_bytes=%d " "(activation=%d sf=%d topk_weights=%d combine=%d shared=%d)", @@ -1801,6 +1849,13 @@ def _region_view( return byte_view.view(dtype).view(shape) def get_regions(self) -> MegaMoeSymmRegions: + inference_mode = torch.is_inference_mode_enabled() + regions = self._regions.get(inference_mode) + if regions is None: + regions = self._regions[inference_mode] = self._build_regions() + return regions + + def _build_regions(self) -> MegaMoeSymmRegions: hidden = self.hidden_size max_t = self.max_tokens_per_rank top_k = self.num_topk @@ -1825,7 +1880,8 @@ def get_regions(self) -> MegaMoeSymmRegions: shared_workspace=self._region_view( "shared_workspace", (self._region_sizes["shared_workspace"],), torch.uint8 ), - peer_offsets=list(self.peer_offsets), + # A tuple: the regions object is shared across launches. + peer_offsets=tuple(self.peer_offsets), rank=self.rank, world_size=self.world_size, ) @@ -1953,6 +2009,17 @@ def set_active_megamoe_cute_weight_cache(cache: Optional[dict]) -> None: global _ACTIVE_MEGAMOE_CUTE_WEIGHT_CACHE _ACTIVE_MEGAMOE_CUTE_WEIGHT_CACHE = cache + # Cute views of the symmetric launch buffers, set by the backend only + # when activation, activation_sf, topk_weights and the combine output are + # its provider's regions and topk_idx is its own staging buffer. The + # provider owns the dict; topk_idx views go to the module's weight cache. + _ACTIVE_MEGAMOE_CUTE_BUFFER_CACHE: Optional[dict] = None + + def set_active_megamoe_cute_buffer_cache(cache: Optional[dict]) -> None: + """Set (``None`` clears) the active provider's cute launch-buffer cache.""" + global _ACTIVE_MEGAMOE_CUTE_BUFFER_CACHE + _ACTIVE_MEGAMOE_CUTE_BUFFER_CACHE = cache + def set_active_megamoe_profiling_scratch(regions) -> None: """Set (``None`` clears) the :class:`MegaMoeSymmRegions` used for AutoTuner profiling on the current call (set by the backend around @@ -2252,7 +2319,14 @@ def unique_id(self): # local_rank is intentionally excluded: every EP rank must run the # SAME tactic and MERGE merges timings by cache key across ranks, # so the key MUST be rank-identical. world_size stays so - # single-rank and multi-rank never share entries. + # single-rank and multi-rank never share entries. The key depends + # only on constructor inputs, so each runner builds it once. + key = self.__dict__.get("_unique_id") + if key is None: + key = self._unique_id = self._build_unique_id() + return key + + def _build_unique_id(self): return _megamoe_problem_key( sm_version=self.sm_version, world_size=self.world_size, @@ -2805,10 +2879,13 @@ def forward( # uint8 workspaces MUST be cute.Pointer, NOT cute.Tensor: the # 32-bit memref shape field overflows once shared_workspace passes # 2 GiB (the kernel addresses by raw base + Int64 byte offset). - _to_cute, _to_cute_ptr, SymBufferHost, cute = _cute_launch_helpers() - # Weights reuse the calling module's cute views; tuning launches convert afresh. - weight_cache = _ACTIVE_MEGAMOE_CUTE_WEIGHT_CACHE - if weight_cache is None or AutoTuner.get().is_tuning_mode: + _to_cute, _, SymBufferHost, cute = _cute_launch_helpers() + # Weights reuse the calling module's cute views and the symmetric + # launch buffers the active provider's; tuning launches convert afresh. + tuning = AutoTuner.get().is_tuning_mode + weight_cache = None if tuning else _ACTIVE_MEGAMOE_CUTE_WEIGHT_CACHE + buffer_cache = None if tuning else _ACTIVE_MEGAMOE_CUTE_BUFFER_CACHE + if weight_cache is None: _to_cute_w = _to_cute else: _to_cute_w = functools.partial(_to_cute_cached, weight_cache) @@ -2817,9 +2894,23 @@ def forward( # present K stride-1 via a transpose VIEW (DLPack carries the # strides); do NOT ``.contiguous()``. output_activation = combine_output.reshape(combine_output.shape[0], self.hidden_size) + if buffer_cache is None: + cute_activation = _to_cute(activation) + cute_activation_sf = _to_cute(activation_sf) + cute_topk_weights = _to_cute(topk_weights) + cute_output = _to_cute(output_activation) + cute_topk_idx = _to_cute(topk_idx) + else: + cute_activation = _to_cute_cached(buffer_cache, activation) + cute_activation_sf = _to_cute_cached(buffer_cache, activation_sf) + cute_topk_weights = _to_cute_cached(buffer_cache, topk_weights) + # Its live rows change per iteration, not between the layers of one. + cute_output = _to_cute_slot(buffer_cache, "output_activation", output_activation) + # topk_idx is the calling module's own staging buffer. + cute_topk_idx = _to_cute_w(topk_idx) runtime_kwargs = dict( - activation=_to_cute(activation), - activation_sf=_to_cute(activation_sf), + activation=cute_activation, + activation_sf=cute_activation_sf, fc1_weight=_to_cute_w(fc1_weight.transpose(1, 2)), fc1_weight_sf=_to_cute_w(fc1_weight_sf), fc2_weight=_to_cute_w(fc2_weight.transpose(1, 2)), @@ -2827,14 +2918,14 @@ def forward( fc1_alpha=_to_cute_w(fc1_alpha, assumed_align=4), fc2_alpha=_to_cute_w(fc2_alpha, assumed_align=4), fc1_norm_const=_to_cute_w(fc1_norm_const, assumed_align=4), - output_activation=_to_cute(output_activation), - local_workspace=_to_cute_ptr(local_workspace), - shared_workspace=_to_cute_ptr(shared_workspace), - stream=cuda.CUstream(torch.cuda.current_stream().cuda_stream), + output_activation=cute_output, + local_workspace=_to_cute_ptr_cached(local_workspace), + shared_workspace=_to_cute_ptr_cached(shared_workspace), + stream=_current_cu_stream(), ) runtime_kwargs.update( - topk_indices=_to_cute(topk_idx), - topk_scores=_to_cute(topk_weights), + topk_indices=cute_topk_idx, + topk_scores=cute_topk_weights, peer_rank_ptr_mapper_host=SymBufferHost( base_address=int(activation.untyped_storage().data_ptr()), offsets=tuple(int(off) for off in peer_offsets), @@ -2846,8 +2937,6 @@ def forward( # Contract: with S > 0 both values must be supplied. Fail loudly # here rather than hand the kernel an unarmed gate. if self.helper_expert_count > 0: - import cutlass - if self._helper_ready_flags is None: raise RuntimeError( "MegaMoE-CuteDSL: helper_expert_count > 0 requires " @@ -2866,7 +2955,7 @@ def forward( f"hold at least world_size={self.world_size} uint64 " f"cells; got numel={self._helper_ready_flags.numel()}." ) - runtime_kwargs["hot_expert_weight_ready_flags"] = _to_cute_ptr( + runtime_kwargs["hot_expert_weight_ready_flags"] = _to_cute_ptr_cached( self._helper_ready_flags ) runtime_kwargs["hot_expert_weight_ready_generation"] = cutlass.Uint64( diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py index f6d6663eb98f..20796c3baa00 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py @@ -357,7 +357,7 @@ class _MegaMoeBuffers: topk_weights: torch.Tensor # (max_T, top_k) float32 combine_output: torch.Tensor # (max_T, top_k, hidden) bf16 shared_workspace: torch.Tensor # (shared_ws_bytes,) uint8 - peer_offsets: List[int] # length == world_size; [0] for single-rank + peer_offsets: Tuple[int, ...] # length == world_size; (0,) for single-rank topk_idx_local: torch.Tensor # (max_T, top_k), int32 on Rubin, otherwise int64 @@ -722,6 +722,8 @@ def __init__( # Per-bucket symmetric providers (bucket -> MegaMoeSymmMemProvider); # filled in ``create_weights`` on the multi-rank EP path. self._symm_providers = {} + # Launch buffers per (bucket, top_k, hidden, device, dtype, inference mode). + self._launch_buffers: Dict[Tuple, _MegaMoeBuffers] = {} self._genphase_direct_input_buckets: set[int] = set() # Lockstep per-chunk cross-rank max for the NEXT run_moe (set via # ``set_adaptive_launch_tokens``); ``None`` -> full bucket. @@ -1186,6 +1188,7 @@ def _alloc_symm_provider(self): top_k = self.routing_method.experts_per_token full_T = self._maxt_buckets[-1] self._profiling_scratch_provider_ref = None + self._launch_buffers.clear() for max_T in self._maxt_buckets: ( shared_workspace_bytes, @@ -1514,20 +1517,16 @@ def run_moe( top_k = int(token_selected_experts.shape[-1]) device = x.device - topk_weights_f32 = token_final_scales.to(torch.float32).contiguous() - # The direct scheduler path borrows contiguous int32 output without a GPU - # operation. Any required conversion is itself a route reader. - if ( - token_selected_experts.dtype != self._topk_idx_dtype - or not token_selected_experts.is_contiguous() - ): - self._wait_rebalance_routes(replica_plan) - topk_idx = token_selected_experts.to(self._topk_idx_dtype).contiguous() + topk_weights_f32 = token_final_scales + if topk_weights_f32.dtype != torch.float32 or not topk_weights_f32.is_contiguous(): + topk_weights_f32 = topk_weights_f32.to(torch.float32).contiguous() + # Routes stay as produced: the staging copy into the kernel's index + # buffer widens int32 to its dtype, and it is the first route reader. return self._run_moe( x=x, x_sf=x_sf, - topk_idx=topk_idx, + topk_idx=token_selected_experts, topk_weights=topk_weights_f32, weight_view=weight_view, num_tokens=num_tokens, @@ -1652,7 +1651,32 @@ def _acquire_buffers( in plain CUDA memory. ``launch_max_T`` selects both the local staging set and (multi-rank) the symmetric provider, which MUST share the same ``max_T`` as the kernel's compile-time ``max_tokens_per_rank``. + Each bundle is built once; views made under inference mode cannot be + updated in place outside it, so the mode is part of the key. """ + key = ( + launch_max_T, + top_k, + hidden, + device, + output_dtype, + torch.is_inference_mode_enabled(), + ) + bufs = self._launch_buffers.get(key) + if bufs is None: + bufs = self._build_buffers( + top_k=top_k, + hidden=hidden, + device=device, + output_dtype=output_dtype, + launch_max_T=launch_max_T, + ) + self._launch_buffers[key] = bufs + return bufs + + def _build_buffers( + self, *, top_k: int, hidden: int, device, output_dtype, launch_max_T: int + ) -> _MegaMoeBuffers: staging = self._ensure_local_staging( top_k=top_k, hidden=hidden, device=device, output_dtype=output_dtype, max_T=launch_max_T ) @@ -1663,7 +1687,7 @@ def _acquire_buffers( topk_weights=staging["topk_weights"], combine_output=staging["combine_output"], shared_workspace=staging["shared_workspace"], - peer_offsets=[0], + peer_offsets=(0,), topk_idx_local=staging["topk_idx"], ) # Multi-rank: the provider must have been rendezvous'd at @@ -1774,6 +1798,7 @@ def _launch_megamoe_kernel( output_dtype: torch.dtype, launch_max_T: int, replica_plan=None, + staged_inputs: bool = False, ) -> torch.Tensor: """Launch the fused MegaMoE CuteDSL kernel and return its reduced output. @@ -1782,7 +1807,9 @@ def _launch_megamoe_kernel( upstream by :meth:`_acquire_buffers`). NVFP4 / FP8-SF dtype views happen through the module-level :func:`_as_nvfp4` / :func:`_as_fp8_sf` helpers (the kernel rejects raw uint8 byte - tensors). + tensors). ``staged_inputs`` states that activation, activation_sf, + topk_idx and topk_weights are this bucket's launch buffers, whose cute + views the op may then reuse. """ # Hand the symmetric profiling scratch (or a deferred factory) to the # op for tuning-mode profiling launches; cleared in the ``finally`` so @@ -1790,6 +1817,7 @@ def _launch_megamoe_kernel( # scratch (ownership: see ``_profiling_scratch_provider_ref``). from tensorrt_llm._torch.moe.custom_ops.cute_dsl_megamoe_custom_op import ( run_cute_dsl_megamoe_nvfp4_blackwell, + set_active_megamoe_cute_buffer_cache, set_active_megamoe_cute_weight_cache, set_active_megamoe_profiling_scratch, set_active_megamoe_profiling_scratch_factory, @@ -1863,6 +1891,9 @@ def _deferred_scratch_factory(_self=self, _get=get_megamoe_profiling_scratch): else run_cute_dsl_megamoe_nvfp4_blackwell ) set_active_megamoe_cute_weight_cache(self._cute_weight_cache) + # The provider owns its symmetric regions, so it owns their cute views. + provider = self._symm_providers.get(launch_max_T) if staged_inputs else None + set_active_megamoe_cute_buffer_cache(provider.cute_views if provider is not None else None) try: megamoe_op( activation=_as_nvfp4(activation), @@ -1915,6 +1946,7 @@ def _deferred_scratch_factory(_self=self, _get=get_megamoe_profiling_scratch): set_active_megamoe_profiling_scratch(None) set_active_megamoe_profiling_scratch_factory(None) set_active_megamoe_cute_weight_cache(None) + set_active_megamoe_cute_buffer_cache(None) if num_tokens == 0: return torch.empty((0, hidden), dtype=output_dtype, device=combine_output.device) # The kernel already collapsed the top-k axis into the (T, 1, hidden) @@ -2158,6 +2190,8 @@ def _run_moe( output_dtype=output_dtype, launch_max_T=launch_max_T, replica_plan=replica_plan, + # Direct routing inputs imply direct activation inputs. + staged_inputs=not use_direct_inputs, ) diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index b53b80cf6253..630acc56a1ec 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -39,8 +39,9 @@ def _extract_class( bases: Tuple[str, ...] = (), ): tree = ast.parse(path.read_text()) + # Breadth-first, so a module-level class wins over a nested one of the same name. original = next( - node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == name + node for node in ast.walk(tree) if isinstance(node, ast.ClassDef) and node.name == name ) if methods is None: selected = original @@ -1587,11 +1588,14 @@ def to_cute(t, assumed_align): conversions.append((t, assumed_align)) return ("cute", len(conversions)) - to_cute_cached = _extract_function( - _TORCH_ROOT / "moe/custom_ops/cute_dsl_megamoe_custom_op.py", - "_to_cute_cached", - {"torch": SimpleNamespace(Tensor=object), "_cute_launch_helpers": lambda: (to_cute,)}, - ) + source = _TORCH_ROOT / "moe/custom_ops/cute_dsl_megamoe_custom_op.py" + namespace = { + "torch": SimpleNamespace(Tensor=object), + "Tuple": Tuple, + "_cute_launch_helpers": lambda: (to_cute,), + } + _extract_function(source, "_cute_view_key", namespace) + to_cute_cached = _extract_function(source, "_to_cute_cached", namespace) def weight(ptr=4096, shape=(8, 16), stride=(16, 1), dtype="fp4", device=0): return SimpleNamespace( @@ -1620,6 +1624,48 @@ def weight(ptr=4096, shape=(8, 16), stride=(16, 1), dtype="fp4", device=0): # The cache belongs to the caller: another module's dict starts empty. assert to_cute_cached({}, weight()) is not first + # A slot keeps only the latest view of a buffer whose live rows change. + to_cute_slot = _extract_function(source, "_to_cute_slot", namespace) + slots = {} + conversions.clear() + rows_8 = to_cute_slot(slots, "output", weight(shape=(8, 16))) + assert to_cute_slot(slots, "output", weight(shape=(8, 16))) is rows_8 + rows_4 = to_cute_slot(slots, "output", weight(shape=(4, 16))) + assert rows_4 is not rows_8 and len(slots) == 1 + assert to_cute_slot(slots, "output", weight(shape=(8, 16))) is not rows_8 + assert len(conversions) == 3 + + +def test_megamoe_cute_pointer_cache_keys_address_and_alignment(): + made = [] + + def to_cute_ptr(t, assumed_align): + made.append((t.data_ptr(), assumed_align)) + return ("ptr", len(made)) + + source = _TORCH_ROOT / "moe/custom_ops/cute_dsl_megamoe_custom_op.py" + cache = {} + namespace = { + "torch": SimpleNamespace(Tensor=object), + "_cute_launch_helpers": lambda: (None, to_cute_ptr), + "_MEGAMOE_CUTE_PTR_CACHE": cache, + "_MEGAMOE_CUTE_PTR_CACHE_MAX_ENTRIES": 2, + } + to_cute_ptr_cached = _extract_function(source, "_to_cute_ptr_cached", namespace) + + def buffer(ptr): + return SimpleNamespace(data_ptr=lambda: ptr) + + first = to_cute_ptr_cached(buffer(4096)) + # Another tensor at the same address gets the same pointer. + assert to_cute_ptr_cached(buffer(4096)) is first + assert to_cute_ptr_cached(buffer(4096), assumed_align=4) is not first + assert len(made) == 2 + # The cache is bounded: a third entry starts it over. + to_cute_ptr_cached(buffer(8192)) + assert len(cache) == 1 + assert to_cute_ptr_cached(buffer(4096)) is not first + def test_route_wait_occurs_after_independent_staging(): source = _TORCH_ROOT / "moe/fused_moe/mega_moe/mega_moe_cute_dsl.py" @@ -1666,6 +1712,71 @@ def run(stage_routes): assert owner._last_staged_T[8] == 8 +def test_megamoe_launch_buffers_are_built_once_per_bucket_and_mode(): + inference = SimpleNamespace(enabled=True) + fake_torch = SimpleNamespace(is_inference_mode_enabled=lambda: inference.enabled) + impl = _extract_class( + _TORCH_ROOT / "moe/fused_moe/mega_moe/mega_moe_cute_dsl.py", + "TrtllmCutedslMegaMoeNvfp4Impl", + {"_acquire_buffers"}, + {"torch": fake_torch}, + ) + built = [] + owner = impl() + owner._launch_buffers = {} + + def build_buffers(**kwargs): + built.append((kwargs["launch_max_T"], inference.enabled)) + return object() + + owner._build_buffers = build_buffers + args = dict(top_k=6, hidden=7168, device="cuda:0", output_dtype="bf16") + + first = owner._acquire_buffers(**args, launch_max_T=8192) + assert owner._acquire_buffers(**args, launch_max_T=8192) is first + assert owner._acquire_buffers(**args, launch_max_T=4096) is not first + # Views made under inference mode are kept apart from normal-mode views. + inference.enabled = False + assert owner._acquire_buffers(**args, launch_max_T=8192) is not first + assert built == [(8192, True), (4096, True), (8192, False)] + + +def test_megamoe_symmetric_regions_are_built_once_per_mode(): + inference = SimpleNamespace(enabled=True) + provider_cls = _extract_class( + _TORCH_ROOT / "moe/custom_ops/cute_dsl_megamoe_custom_op.py", + "MegaMoeSymmMemProvider", + {"get_regions"}, + {"torch": SimpleNamespace(is_inference_mode_enabled=lambda: inference.enabled)}, + ) + provider = provider_cls() + provider._regions = {} + built = [] + provider._build_regions = lambda: built.append(inference.enabled) or object() + + regions = provider.get_regions() + assert provider.get_regions() is regions + inference.enabled = False + assert provider.get_regions() is not regions + assert built == [True, False] + + +def test_megamoe_runner_identity_is_built_once(): + runner_cls = _extract_class( + _TORCH_ROOT / "moe/custom_ops/cute_dsl_megamoe_custom_op.py", + "Sm100MegaMoENvfp4Runner", + {"unique_id"}, + {}, + ) + runner = runner_cls() + builds = [] + runner._build_unique_id = lambda: builds.append(1) or (100, 8, 6) + + assert runner.unique_id() == (100, 8, 6) + assert runner.unique_id() == (100, 8, 6) + assert builds == [1] + + def test_warmup_owner_handoff_uses_deduplicated_manager_public_api(): events = [] main = SimpleNamespace(cuda_stream=11, priority=0, device=0) diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index 27486ffec9c2..4734bb8312b5 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -2211,6 +2211,92 @@ def test_megamoe_replica_plan_supplies_launch_facts_and_route_wait(): wait_for_routes.assert_called_once_with() +@pytest.mark.cpu_only +def test_megamoe_staging_widens_int32_routes_after_the_route_wait(): + moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) + torch.nn.Module.__init__(moe) + moe._last_staged_T = {8: 6} + trace = [] + moe._wait_rebalance_routes = lambda replica_plan=None: trace.append(("wait", replica_plan)) + bufs = SimpleNamespace( + activation=torch.zeros(8, 16, dtype=torch.uint8), + activation_sf=torch.zeros(8, 4, dtype=torch.uint8), + topk_weights=torch.ones(8, 4, dtype=torch.float32), + topk_idx_local=torch.full((8, 4), 5, dtype=torch.int64), + ) + routes = torch.tensor([[0, 3, -1, 2**31 - 1], [7, 6, 5, 4], [1, 1, 1, 1]], dtype=torch.int32) + weights = torch.full((3, 4), 0.25, dtype=torch.float32) + + moe._stage_inputs( + bufs=bufs, + x=torch.ones(3, 16, dtype=torch.uint8), + x_sf=torch.ones(3, 4, dtype=torch.uint8), + topk_idx=routes, + topk_weights=weights, + num_tokens=3, + top_k=4, + stage_activation=True, + replica_plan="plan", + ) + + # One copy widens the int32 routes into the int64 kernel index buffer. + assert bufs.topk_idx_local.dtype == torch.int64 + assert torch.equal(bufs.topk_idx_local[:3], routes.to(torch.int64)) + # Rows a previous launch wrote, and only those, are reset to -1. + assert torch.equal(bufs.topk_idx_local[3:6], torch.full((3, 4), -1, dtype=torch.int64)) + assert torch.equal(bufs.topk_idx_local[6:], torch.full((2, 4), 5, dtype=torch.int64)) + assert torch.equal(bufs.topk_weights[:3], weights) + assert not bufs.topk_weights[3:].any() + assert trace == [("wait", "plan")] + + +@pytest.mark.cpu_only +def test_megamoe_run_moe_hands_routes_to_staging_unconverted(): + moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) + torch.nn.Module.__init__(moe) + moe.ep_size = 1 + moe._weights_created = True + moe.quant_method = object() + moe.hidden_size = 16 + moe.dtype = torch.bfloat16 + moe._topk_idx_dtype = torch.int64 + moe._build_weight_view = lambda: "weights" + moe._wait_rebalance_routes = MagicMock() + captured = {} + moe._run_moe = lambda **kwargs: captured.update(kwargs) or "out" + routes = torch.zeros(3, 4, dtype=torch.int32) + scales = torch.ones(3, 4, dtype=torch.float32) + + out = moe.run_moe( + MoERunContext( + token_selected_experts=routes, + token_final_scales=scales, + x=torch.zeros(3, 8, dtype=torch.uint8), + x_sf=torch.zeros(3, 4, dtype=torch.uint8), + output_dtype=torch.bfloat16, + ) + ) + + assert out == "out" + # No conversion kernel and no early route wait: staging reads the routes. + assert captured["topk_idx"] is routes + assert captured["topk_weights"] is scales + moe._wait_rebalance_routes.assert_not_called() + + half_scales = scales.to(torch.float16) + moe.run_moe( + MoERunContext( + token_selected_experts=routes, + token_final_scales=half_scales, + x=torch.zeros(3, 8, dtype=torch.uint8), + x_sf=torch.zeros(3, 4, dtype=torch.uint8), + output_dtype=torch.bfloat16, + ) + ) + assert captured["topk_weights"].dtype == torch.float32 + assert torch.equal(captured["topk_weights"], scales) + + @pytest.mark.cpu_only def test_megamoe_replica_plan_rejects_unallocated_geometry(): moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) From dfefce52339358b922479cbeb84176138db1d32a Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 14:55:06 -0700 Subject: [PATCH 25/59] [None][perf] Skip no-op slices, splits and casts in the fused-comm MoE scheduler FusedCommMoEScheduler sliced x, the router logits and the input ids to the local token count, split them into chunks and cast the routes and scales on every layer, even when the slice covers the whole tensor, a single chunk covers it, or the dtype already matches. In context-only serving each of these is a no-op that still costs a dispatcher call per layer. Use the tensor itself in those cases; padded inputs and multi-chunk forwards are sliced and split as before. Signed-off-by: jintaop --- .../_torch/moe/fused_moe/moe_scheduler.py | 47 +++++++++++++------ tests/unittest/_torch/moe/test_moe_backend.py | 47 +++++++++++++++++++ 2 files changed, 80 insertions(+), 14 deletions(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py index c5708c1cdac8..a01537dda4f6 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py @@ -71,6 +71,18 @@ from .configurable_moe import ConfigurableMoE +def _leading_rows(t: torch.Tensor, rows: int) -> torch.Tensor: + """Return ``t[:rows]``, or ``t`` itself when it already has ``rows`` rows.""" + return t if t.shape[0] == rows else t[:rows] + + +def _split_rows(t: torch.Tensor, sizes: List[int]) -> List[torch.Tensor]: + """Return ``list(t.split(sizes))``; one chunk covering ``t`` is ``t`` itself.""" + if len(sizes) == 1 and sizes[0] == t.shape[0]: + return [t] + return list(t.split(sizes)) + + def _wait_input_ready(event: Optional[torch.cuda.Event]) -> None: """Join a pending ``input_ready_event`` before the first read of ``x``. @@ -1234,12 +1246,12 @@ def forward( if input_ids_real is not None: chunk_size_list = all_rank_chunk_size_list[ep_rank] input_ids_chunks = ( - list(input_ids_real.split(chunk_size_list)) if input_ids_real.numel() > 0 else [] + _split_rows(input_ids_real, chunk_size_list) if input_ids_real.numel() > 0 else [] ) routing_aux_chunks = [] if routing_aux is not None and x_real.shape[0] > 0: - routing_aux_chunks = list( - routing_aux[: x_real.shape[0]].split(all_rank_chunk_size_list[ep_rank]) + routing_aux_chunks = _split_rows( + _leading_rows(routing_aux, x_real.shape[0]), all_rank_chunk_size_list[ep_rank] ) outputs = self._run_chunks( x_chunks, @@ -1326,9 +1338,9 @@ def _strip_adp_padding( ) # ADP padding stripped before split, else trailing rows silently # drift into chunk-0 or torch.split shape-errors. - x_real = x[:real_local] - rl_real = router_logits[:real_local] - input_ids_real = input_ids[:real_local] if input_ids is not None else None + x_real = _leading_rows(x, real_local) + rl_real = _leading_rows(router_logits, real_local) + input_ids_real = _leading_rows(input_ids, real_local) if input_ids is not None else None return x_real, rl_real, input_ids_real, real_all_rank_num_tokens, ep_rank, had_meta def _compute_chunk_layout( @@ -1368,8 +1380,8 @@ def _compute_chunk_layout( # ``else []`` shortcut for real_local == 0: equivalent to # x_real.split([0]*num_chunks) but skips the no-op torch call. The # zero-token fallback in ``_run_chunks`` then fires for every chunk. - x_chunks = list(x_real.split(chunk_size_list)) if real_local > 0 else [] - rl_chunks = list(rl_real.split(chunk_size_list)) if real_local > 0 else [] + x_chunks = _split_rows(x_real, chunk_size_list) if real_local > 0 else [] + rl_chunks = _split_rows(rl_real, chunk_size_list) if real_local > 0 else [] return num_chunks, x_chunks, rl_chunks, all_rank_chunk_size_list def _run_chunks( @@ -1490,9 +1502,11 @@ def _forward_chunk( num_tokens = x.shape[0] assert num_tokens <= x.shape[0], f"num_tokens ({num_tokens}) > x.shape[0] ({x.shape[0]})" - x_chunk_real = x[:num_tokens] - router_logits_chunk_real = router_logits[:num_tokens] - input_ids_chunk_real = input_ids[:num_tokens] if input_ids is not None else None + x_chunk_real = _leading_rows(x, num_tokens) + router_logits_chunk_real = _leading_rows(router_logits, num_tokens) + input_ids_chunk_real = ( + _leading_rows(input_ids, num_tokens) if input_ids is not None else None + ) # ----- EPLB: drain previous CPU rebalance ----- # Static EPLB early-returns inside the helper; only the dynamic @@ -1509,15 +1523,20 @@ def _forward_chunk( ) else: token_selected_experts, token_final_scales = moe.routing_method.apply_with_aux( - router_logits_chunk_real, input_ids_chunk_real, routing_aux[:num_tokens] + router_logits_chunk_real, + input_ids_chunk_real, + _leading_rows(routing_aux, num_tokens), ) + if token_selected_experts.dtype != torch.int32: + token_selected_experts = token_selected_experts.to(torch.int32) token_selected_experts = moe._maybe_get_perfect_router_expert_ids( - token_selected_experts.to(torch.int32) + token_selected_experts ) route_capture = get_active_route_capture() if route_capture is not None: # R3: in-graph device-buffer capture route_capture.capture(moe.layer_idx, token_selected_experts, row_offset) - token_final_scales = token_final_scales.to(torch.float32) + if token_final_scales.dtype != torch.float32: + token_final_scales = token_final_scales.to(torch.float32) else: device = x.device token_selected_experts = torch.empty( diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index 4734bb8312b5..68cf1156ec07 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -373,6 +373,53 @@ def test_fused_comm_scheduler_indexes_full_ep_token_metadata(): assert had_meta +def _split_chunk(tokens, chunks): + return [tokens // chunks + 1] * (tokens % chunks) + [tokens // chunks] * ( + chunks - tokens % chunks + ) + + +@pytest.mark.cpu_only +def test_fused_comm_scheduler_passes_unpadded_single_chunk_inputs_through(): + moe = SimpleNamespace( + mapping=SimpleNamespace(moe_ep_rank=1, moe_ep_size=2), + moe_max_num_tokens=8, + split_chunk=_split_chunk, + ) + scheduler = FusedCommMoEScheduler(moe) + x = torch.randn(6, 8) + router_logits = torch.randn(6, 4) + input_ids = torch.arange(6) + + x_real, rl_real, ids_real, counts, ep_rank, had_meta = scheduler._strip_adp_padding( + x, router_logits, input_ids, [4, 6] + ) + # This rank carries no ADP padding and fits one chunk: no view is taken. + assert x_real is x and rl_real is router_logits and ids_real is input_ids + assert (counts, ep_rank, had_meta) == ([4, 6], 1, True) + num_chunks, x_chunks, rl_chunks, chunk_sizes = scheduler._compute_chunk_layout( + x_real, rl_real, counts, ep_rank + ) + assert num_chunks == 1 + assert x_chunks[0] is x and rl_chunks[0] is router_logits + assert chunk_sizes == [[4], [6]] + + # Padding is still stripped and several chunks are still split. + x_real, rl_real, _, counts, ep_rank, _ = scheduler._strip_adp_padding( + x, router_logits, None, [4, 5] + ) + assert x_real.shape[0] == 5 and torch.equal(x_real, x[:5]) + moe.moe_max_num_tokens = 2 + num_chunks, x_chunks, rl_chunks, chunk_sizes = scheduler._compute_chunk_layout( + x_real, rl_real, counts, ep_rank + ) + assert num_chunks == 3 + assert [chunk.shape[0] for chunk in x_chunks] == [2, 2, 1] + assert torch.equal(torch.cat(x_chunks), x[:5]) + assert torch.equal(torch.cat(rl_chunks), router_logits[:5]) + assert chunk_sizes == [[2, 1, 1], [2, 2, 1]] + + def test_fused_comm_scheduler_rejects_partial_ep_token_metadata(): mapping = SimpleNamespace(moe_ep_rank=1, moe_ep_size=4) scheduler = FusedCommMoEScheduler(SimpleNamespace(mapping=mapping)) From 19d8209f85a72f1f080b5b2d7862fde5e28a8658 Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 14:56:48 -0700 Subject: [PATCH 26/59] [None][perf] Reuse autotune-opted MegaMoE runners outside tuning mode With MEGAMOE_TACTIC_AUTOTUNE=1 every MegaMoE launch constructed a new runner and a new TuningConfig, because a runner that may profile owns profiling scratch. Outside tuning mode choose_one only looks up the tuned tactic and never profiles, so cache that runner like the opted-out one and create fresh runners only while the AutoTuner is tuning. Each runner builds its TuningConfig once; the profiling pre-hook stays bound to that runner. The autotune path now also drops its READY-flag reference after the launch, as the opted-out path does. Signed-off-by: jintaop --- .../custom_ops/cute_dsl_megamoe_custom_op.py | 38 ++++++++------ .../_torch/moe/test_dynamic_eplb_contracts.py | 50 +++++++++++++++++++ 2 files changed, 73 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py b/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py index fc4ea0aed7be..ddc8cd0e371c 100644 --- a/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py +++ b/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py @@ -1522,8 +1522,9 @@ def enumerate_megamoe_candidate_tactics( # freed when the next one allocates (only the winner is reused). _MEGAMOE_TUNING_WORKSPACE_KEYS: set = set() - # Cache runners by constructor inputs only outside tactic autotuning. - # Reset profiling scratch on reuse and overwrite READY inputs for every call; + # Cache runners by constructor inputs, except while the AutoTuner profiles + # an autotune opt-in. Reset profiling scratch on reuse and overwrite READY + # inputs for every call. _MEGAMOE_RUNNER_CACHE: dict = {} @functools.lru_cache(maxsize=1) @@ -2488,13 +2489,15 @@ def get_tuning_config(self) -> TuningConfig: autotuner does not double-enumerate tile sizes for independent token axes. - Rebuilt on every call, NEVER cached across runners: + Built once per runner and NEVER shared across runners: ``inputs_pre_hook`` is bound to THIS runner's ``_profiling_scratch``; - a cached hook bound to a dead runner would mix non-symmetric - profiling tensors with the live runner's peer offsets (IMA + hang). + a hook bound to another runner would mix non-symmetric profiling + tensors with the live runner's peer offsets (IMA + hang). """ - - config = TuningConfig( + config = self.__dict__.get("_tuning_config") + if config is not None: + return config + config = self._tuning_config = TuningConfig( dynamic_tensor_specs=( DynamicTensorSpec( 0, 0, get_last_power_of_2_num_tokens_buckets, last_positive_power_of_2 @@ -3041,8 +3044,10 @@ def _make(): helper_expert_count=helper_expert_count, ) - # Autotuning runners own profiling scratch and are not reused. - if tactic_autotune: + # A runner that may profile owns profiling scratch, so tuning-mode calls + # with the autotune opt-in get a fresh one. Outside tuning mode + # choose_one only looks up the tuned tactic, and the runner is reused. + if tactic_autotune and AutoTuner.get().is_tuning_mode: return _make() runner = _MEGAMOE_RUNNER_CACHE.get(args) if runner is None: @@ -3243,12 +3248,15 @@ def run_cute_dsl_megamoe_nvfp4_blackwell( runner._profiling_scratch = None runner._profiling_scratch_factory = None # Real run always uses the caller's staging buffer + peer_offsets. - runner( - inputs, - tactic=best_tactic, - peer_offsets=peer_offsets, - shared_workspace=shared_workspace, - ) + try: + runner( + inputs, + tactic=best_tactic, + peer_offsets=peer_offsets, + shared_workspace=shared_workspace, + ) + finally: + runner._helper_ready_flags = None # Eager callers may call ``run_cute_dsl_megamoe_nvfp4_blackwell`` directly # and skip the custom-op dispatch; traced code must use the registered op. diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index 630acc56a1ec..a33c644e446a 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -1761,6 +1761,39 @@ def test_megamoe_symmetric_regions_are_built_once_per_mode(): assert built == [True, False] +def test_megamoe_tuning_config_is_built_once_per_runner(): + configs = [] + + def tuning_config(**kwargs): + configs.append(kwargs) + return SimpleNamespace(**kwargs) + + runner_cls = _extract_class( + _TORCH_ROOT / "moe/custom_ops/cute_dsl_megamoe_custom_op.py", + "Sm100MegaMoENvfp4Runner", + {"get_tuning_config", "_autotuner_inputs_pre_hook"}, + { + "TuningConfig": tuning_config, + "DynamicTensorSpec": lambda *args: args, + "ConstraintSpec": lambda *args: args, + "DistributedTuningStrategy": SimpleNamespace(MERGE="merge"), + "get_last_power_of_2_num_tokens_buckets": None, + "last_positive_power_of_2": None, + "_megamoe_autotune_num_tokens": None, + }, + ) + first, second = runner_cls(), runner_cls() + for runner in (first, second): + runner.max_tokens_per_rank = 8192 + + config = first.get_tuning_config() + assert first.get_tuning_config() is config + # The profiling pre-hook stays bound to the runner that built the config. + assert config.inputs_pre_hook.__self__ is first + assert second.get_tuning_config().inputs_pre_hook.__self__ is second + assert len(configs) == 2 + + def test_megamoe_runner_identity_is_built_once(): runner_cls = _extract_class( _TORCH_ROOT / "moe/custom_ops/cute_dsl_megamoe_custom_op.py", @@ -1942,7 +1975,9 @@ def __init__(self, **kwargs): self._profiling_scratch_factory = object() created.append(self) + tuner = SimpleNamespace(is_tuning_mode=False) namespace = { + "AutoTuner": SimpleNamespace(get=lambda: tuner), "Sm100MegaMoENvfp4Runner": FakeRunner, "_MEGAMOE_RUNNER_CACHE": {}, } @@ -1983,11 +2018,26 @@ def __init__(self, **kwargs): assert production._profiling_scratch is None assert production._profiling_scratch_factory is None + # While the AutoTuner profiles, an autotune opt-in gets a fresh runner + # that owns its profiling scratch. + tuner.is_tuning_mode = True first_tuning = get_runner(**kwargs, tactic_autotune=True) second_tuning = get_runner(**kwargs, tactic_autotune=True) assert first_tuning is not second_tuning + assert first_tuning is not production assert len(created) == 3 + # Outside tuning mode it only looks up its tactic: one runner is reused, + # separate from the opted-out production runner. + tuner.is_tuning_mode = False + serving = get_runner(**kwargs, tactic_autotune=True) + serving._profiling_scratch = "stale" + assert get_runner(**kwargs, tactic_autotune=True) is serving + assert serving is not production + assert serving._profiling_scratch is None + assert get_runner(**kwargs, tactic_autotune=False) is production + assert len(created) == 4 + def test_on_autotune_uses_distinct_tactics(): helpers = _autotune_helpers(None) From a79e5993658c43eac4fe8a6bea1214d51806ce22 Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 15:08:00 -0700 Subject: [PATCH 27/59] [None][perf] Refuse new dynamic EPLB plans while launched routes are unwaited The scheduler group holds the routes of a HALO-Q launch until MAIN waits for its route epoch. Take that reference only right before the launch, so a plan that fails earlier leaves nothing behind, and refuse the next plan while a launch's routes are still held: only a submission that failed after launching leaves them, and dropping them would let the allocator hand their block to MAIN work while HALO-Q may still read it. Signed-off-by: jintaop --- .../mega_moe/rebalance_slot_scheduler.py | 5 ++++- .../_torch/moe/test_dynamic_eplb_contracts.py | 16 ++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py index 7d1d19d90e83..693f3b042e71 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py @@ -405,6 +405,9 @@ def plan_schedule(self, logical_expert_ids: torch.Tensor) -> Tuple[Any, int]: self._check_owner() if self._plan_part is not None or self._generation != self._finished_generation: raise RuntimeError("Previous rebalance plan has no paired finish") + if self._routes_in_flight is not None: + # Only a failed submission leaves routes that MAIN never waited for. + raise RuntimeError("Routes of an earlier HALO-Q launch were never ordered on MAIN") if ( not isinstance(logical_expert_ids, torch.Tensor) or logical_expert_ids.ndim != 2 @@ -420,7 +423,6 @@ def plan_schedule(self, logical_expert_ids: torch.Tensor) -> Tuple[Any, int]: # Both waits capture their event's current record before a later plan # can re-record it. No host synchronization or temporary event is needed. - self._routes_in_flight = logical_expert_ids self._record_event(self._input_ready_handle, self._execution_stream_handle) self._wait_event(self._stream_handle, self._input_ready_handle) outputs = self.scheduler.outputs @@ -429,6 +431,7 @@ def plan_schedule(self, logical_expert_ids: torch.Tensor) -> Tuple[Any, int]: self.lease.prove(self._pending_release[0]) self._pending_release = None route_target = self._route_target + self.auxiliary_sms + self._routes_in_flight = logical_expert_ids ticket = self.broadcaster.submit_fused( outputs, logical_expert_ids, diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index a33c644e446a..426644b96faa 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -1229,6 +1229,7 @@ def test_fused_submit_orders_generations_and_defers_the_route_wait(auxiliary_sms setattr(group, name, 0) group._plan_part = None group._pending_release = None + group._routes_in_flight = None group._closed = False group._owner_thread_id = None group._execution_stream_handle = None @@ -1335,6 +1336,21 @@ def submit_fused(result, routes, epoch, target, *, spin_cycles): group.plan_schedule(_Routes(8, trace)) assert not any(entry[0] == "FUSED" for entry in trace) assert group._plan_part is None + # Nothing was launched, so the routes were never handed to HALO-Q. + assert group._routes_in_flight is None + + # A submission that fails after the launch keeps its routes alive, and the + # group refuses further plans rather than let the allocator reuse them. + group._driver.cuStreamWaitEvent = lambda *args: (0,) + group._pending_release = None + group._spin_bound = SimpleNamespace(cycles=lambda: _SERVING_SPIN_CYCLES) + group.broadcaster.submit_fused = lambda *args, **kwargs: SimpleNamespace(generation=99) + stuck = _Routes(8, trace) + with pytest.raises(RuntimeError, match="generation mismatch"): + group.plan_schedule(stuck) + assert group._routes_in_flight is stuck + with pytest.raises(RuntimeError, match="never ordered on MAIN"): + group.plan_schedule(_Routes(8, trace)) # A closed group rejects every route wait, including one of a waited generation. group._closed = True From 489e5551ab492ec0a916eb6af2e8d903b13ae0af Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 14:47:07 -0700 Subject: [PATCH 28/59] [None][perf] Use the fused shared-expert kernel on the whole SM100 family With per-layer dynamic EPLB, DeepSeek-V4 runs its shared expert while the fused HALO-Q + in-switch copy kernel plans and copies helper weights. The fused shared FC1/activation/FC2 kernel takes an SM budget that leaves the copy kernel its auxiliary SMs. Select it on every SM100-family GPU instead of SM107 only, so GB200 and GB300 no longer run the shared expert's two DeepGEMM GEMMs alongside the copy. Signed-off-by: jintaop --- examples/wide_ep/ep_load_balancer/README.md | 2 +- tensorrt_llm/_torch/models/modeling_deepseekv4.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/wide_ep/ep_load_balancer/README.md b/examples/wide_ep/ep_load_balancer/README.md index 0ffb784680fa..e8d8644334cf 100644 --- a/examples/wide_ep/ep_load_balancer/README.md +++ b/examples/wide_ep/ep_load_balancer/README.md @@ -309,7 +309,7 @@ Per-layer dynamic EPLB has the following constraints: * **Parallelism**: attention DP (`enable_attention_dp: true`) with pure expert parallelism: `moe_expert_parallel_size` equals the number of GPUs, with no MoE tensor parallelism and no pipeline parallelism. The EP size must be even and between 2 and 32, and the MoE smart router (`moe_cluster_size` greater than 1) is not supported. Unlike `standard` mode, an unsupported model, an EP size of 1, or the smart router raises an error instead of silently disabling the load balancer. * **CUDA graphs**: not supported. A non-null `cuda_graph_config` is replaced with `null` and a warning is logged, and any `prefill_cuda_graph_backend` other than `DISABLED` is rejected. * **Launch queues**: CUDA must be initialized with `CUDA_SCALE_LAUNCH_QUEUES=4x`. The LLM API sets it automatically, in the main process and in every worker, before CUDA is initialized. If the process has already initialized CUDA when the LLM is created, construction fails; export `CUDA_SCALE_LAUNCH_QUEUES=4x` before starting such a process. -* **`auxiliary_sms`**: the number of SMs of the kernel that runs the HALO-Q scheduler and then the copy, one CTA per SM. It defaults to 8, must be between 1 and 8, and must be lower than the GPU's SM count. Input quantization and, on SM107, DeepSeek-V4's fused shared-expert kernel leave these SMs free, so planning and copying run alongside them. +* **`auxiliary_sms`**: the number of SMs of the kernel that runs the HALO-Q scheduler and then the copy, one CTA per SM. It defaults to 8, must be between 1 and 8, and must be lower than the GPU's SM count. Input quantization and, on the SM100 family, DeepSeek-V4's fused shared-expert kernel leave these SMs free, so planning and copying run alongside them. * **`helper_slot_mode`**: `per_layer_resident` (the default) gives every MoE layer its own helper slots and skips copying experts that they still hold from the previous forward pass. It costs one helper set per MoE layer, about 6.4 GiB per GPU for DeepSeek-V4-Pro with 3 helper slots per rank. `shared_banks` has all MoE layers alternate between two shared helper banks and copies every helper expert in every forward pass. * **DWDP**: not supported; `dwdp_config` must keep `dwdp_size` at 1. * **Fields**: only `mode`, `num_slots`, `auxiliary_sms` and `helper_slot_mode`. `num_slots` is required, must be greater than the expert count, and must be divisible by the EP size. `layer_updates_per_iter` and `initial_global_assignments` are `standard`-only and rejected in this mode. diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index e7f86ab7846c..9139c9f3835c 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -47,7 +47,7 @@ import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils from tensorrt_llm._ipc_utils import can_access_peer from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict -from tensorrt_llm._utils import get_sm_version, is_sm_120f +from tensorrt_llm._utils import get_sm_version, is_sm_100f, is_sm_120f from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig @@ -1736,7 +1736,7 @@ def __init__( use_fused_fc12 = ( self._per_layer_dynamic_eplb_enabled and self._shared_fc12_backend_supported(rebalance_backend) - and get_sm_version() == 107 + and is_sm_100f() and self._shared_fc12_quantization_supported(gate_up_quant_config, down_quant_config) and not model_config.use_cuda_graph ) From 64262a78f50de7eed26e3f8e16f6b423e83c2758 Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 16:30:07 -0700 Subject: [PATCH 29/59] [None][perf] Keep DeepSeek-V4's shared-expert kernel independent of EPLB DeepSeek-V4 selected the fused shared FC1/activation/FC2 kernel only with per-layer dynamic EPLB, so turning the load balancer on also swapped the shared expert from GatedMLP's two DeepGEMM GEMMs to the fused kernel. Select the fused kernel whenever the routed MegaMoE backend, the final projection quantization and the SM100 family allow it and the model runs eagerly, with or without the load balancer. An eager EPLB-off baseline then runs the same shared expert, and enabling EPLB adds only expert placement and the fused kernel's SM budget. Signed-off-by: jintaop --- tensorrt_llm/_torch/models/modeling_deepseekv4.py | 10 +++++----- .../_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md | 13 ++++++++----- 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 9139c9f3835c..4d6c29ba6b9c 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -1730,12 +1730,12 @@ def __init__( int(load_balancer_config.auxiliary_sms) if self._per_layer_dynamic_eplb_enabled else 0 ) rebalance_backend = getattr(self.experts, "backend", self.experts) - # The fused shared kernel is part of the opt-in per-layer dynamic path. Ask - # the selected routed backend for the capability instead of naming a - # concrete implementation here. + # The fused shared FC1/activation/FC2 kernel comes with the routed MegaMoE + # backend; ask that backend for the capability instead of naming a concrete + # implementation here. Per-layer dynamic EPLB is not a condition, but the + # kernel runs only in eager mode, which that load balancer requires. use_fused_fc12 = ( - self._per_layer_dynamic_eplb_enabled - and self._shared_fc12_backend_supported(rebalance_backend) + self._shared_fc12_backend_supported(rebalance_backend) and is_sm_100f() and self._shared_fc12_quantization_supported(gate_up_quant_config, down_quant_config) and not model_config.use_cuda_graph diff --git a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md index 8b38336449c7..9bab613bc800 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -911,11 +911,14 @@ set wrong has a test that catches it: handle's `allocate_arena` and exposes `live_weight_arena`, `ep_process_group`, and `max_tokens_per_rank_capacity` to `bind_backend` instead of private attributes. -- Per-layer dynamic EPLB selects the fused shared FC12 kernel only when the routed - backend declares per-layer dynamic support and the final projection - quantization, hardware, and graph mode support that kernel. Other runs keep - the model's existing shared-MLP backend. The active configuration supplies - the optional SM budget only to the fused kernel. +- DeepSeek-V4 selects the fused shared FC12 kernel when the routed backend + declares per-layer dynamic support, the final projection quantization and + hardware (SM100 family) support that kernel, and the model runs eagerly + (`cuda_graph_config: null`). Per-layer dynamic EPLB is not a condition, but it + requires eager mode, so an EPLB-off baseline runs the same shared kernel only + with `cuda_graph_config: null`. Other runs keep the model's existing shared-MLP + backend. With per-layer dynamic EPLB, the active configuration supplies the + fused kernel an SM budget that leaves the copy kernel its auxiliary SMs. ## Canonical Examples From e6219459523f7c2c09bfd684d78482ad0413732e Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 16:30:20 -0700 Subject: [PATCH 30/59] [None][fix] Round the fused shared-expert block scales to UE8M0 at load on SM100 and SM103 Linear converts FP8 block scales to the CuTe DSL MXFP8 layout (UE8M0 K32 R128c4) only on SM107, where its own GEMM reads that layout. On SM100 and SM103 it keeps FP32 K128 scales, so the fused shared FC12 kernel rejected them at its first forward. MegaMoESharedMLP now asks its two projections to round those scales to UE8M0 during Linear's post-load transform, as the DeepGEMM transform does, keeping the FP32 K128 layout so Linear's own GEMM stays valid. The fused kernel derives its R128c4 copies from them and never writes to a parameter at forward time. Signed-off-by: jintaop --- tensorrt_llm/_torch/modules/linear.py | 14 ++- .../_torch/modules/megamoe_shared_mlp.py | 17 ++- .../test_lists/test-db/l0_b200.yml | 1 + .../modules/test_megamoe_shared_mlp_scales.py | 101 ++++++++++++++++++ 4 files changed, 130 insertions(+), 3 deletions(-) create mode 100644 tests/unittest/_torch/modules/test_megamoe_shared_mlp_scales.py diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index a4818fd02a71..ea9419f18bae 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -1642,7 +1642,10 @@ def transform_weights(self, module: Linear) -> None: or module.disable_deep_gemm)) or is_sm_120f() use_indexer_q_cutedsl_layout = (use_deep_gemm_layout and getattr( module, "use_indexer_q_cutedsl_fusion", False)) - if use_deep_gemm_layout or use_indexer_q_cutedsl_layout: + # Consumers that read UE8M0 scales from the FP32 K128 layout (e.g. the + # fused shared FC12 kernel) ask for the rounding alone. + round_to_ue8m0 = getattr(module, "round_block_scales_to_ue8m0", False) + if use_deep_gemm_layout or use_indexer_q_cutedsl_layout or round_to_ue8m0: weight, weight_scale = resmooth_to_fp8_e8m0(module.weight, module.weight_scale) @@ -1688,6 +1691,15 @@ def transform_weights(self, module: Linear) -> None: if get_sm_version() != 107: module.fp8_low_m_gemm = Fp8BlockScaleLowMGemmRange.for_weight( *weight.shape) + elif round_to_ue8m0: + # Same FP32 K128 layout; the rounded tensors may live on another device. + replace_parameter_and_save_metadata( + module, "weight", nn.Parameter(weight, requires_grad=False), + module.rebuild_tensor_metadata) + replace_parameter_and_save_metadata( + module, "weight_scale", + nn.Parameter(weight_scale, requires_grad=False), + module.rebuild_tensor_metadata) class NVFP4LinearMethod(LinearMethodBase): diff --git a/tensorrt_llm/_torch/modules/megamoe_shared_mlp.py b/tensorrt_llm/_torch/modules/megamoe_shared_mlp.py index 4aea1efc61d8..0e0d416c437e 100644 --- a/tensorrt_llm/_torch/modules/megamoe_shared_mlp.py +++ b/tensorrt_llm/_torch/modules/megamoe_shared_mlp.py @@ -20,6 +20,8 @@ import torch import torch.nn.functional as F +from tensorrt_llm.quantization.utils.fp8_utils import transform_k128_scales_to_cutedsl_mxfp8_layout + from ..distributed import AllReduceParams from ..model_config import ModelConfig from .gated_mlp import GatedMLP @@ -116,8 +118,8 @@ def __init__( overridden_tp_size=overridden_tp_size, reduce_output=False, layer_idx=layer_idx, - # This selects the existing Linear post-load UE8M0 conversion, including when - # the caller's dense GEMM preference is unset. + # On SM107 this selects Linear's post-load CuTe DSL MXFP8 conversion. On + # SM100/103 Linear then keeps FP32 K128 scales instead of DeepGEMM's layout. use_cute_dsl_blockscaling_mm=True, disable_deep_gemm=disable_deep_gemm, is_shared_expert=is_shared_expert, @@ -126,6 +128,9 @@ def __init__( swiglu_beta=swiglu_beta, ) self._validate_fp8_quantization() + # The fused kernel reads UE8M0 scales; Linear rounds FP32 block scales at load. + for projection in (self.gate_up_proj, self.down_proj): + projection.round_block_scales_to_ue8m0 = True self._fc12_reserved_sms = fc12_reserved_sms self._fc12_max_tokens = int(config.max_num_tokens) if self._fc12_max_tokens <= 0: @@ -198,6 +203,14 @@ def prepare_fc12_weights(self) -> None: ) if any(t.device != fc1.device for t in sources): raise ValueError("Shared FC12 weights and scales must be on the same device") + if fc1_sf.dtype == torch.float32 and fc2_sf.dtype == torch.float32: + # Linear kept UE8M0-valued FP32 K128 scales (SM100/103). + fc1_sf = transform_k128_scales_to_cutedsl_mxfp8_layout( + fc1_sf, mn=fc1.shape[0], k=fc1.shape[1] + ) + fc2_sf = transform_k128_scales_to_cutedsl_mxfp8_layout( + fc2_sf, mn=fc2.shape[0], k=fc2.shape[1] + ) logical_fc1_sf = _unpack_r128c4(fc1_sf, 2 * intermediate, self.hidden_size // 32) _unpack_r128c4(fc2_sf, self.hidden_size, intermediate // 32) # Byte views keep packing independent of PyTorch FP8 indexing support. diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index e8e1f903dea9..a5bcb895a23c 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -142,6 +142,7 @@ l0_b200: - unittest/_torch/modules/test_engram.py - unittest/_torch/modules/test_gate_up_swiglu_quack.py - unittest/_torch/modules/test_multi_stream_utils.py + - unittest/_torch/modules/test_megamoe_shared_mlp_scales.py # ------------- DeepSeek-V4.1 model, Engram and DSpark --------------- - unittest/_torch/modeling/test_modeling_deepseekv41.py - unittest/_torch/modules/test_engram_decode_hash.py diff --git a/tests/unittest/_torch/modules/test_megamoe_shared_mlp_scales.py b/tests/unittest/_torch/modules/test_megamoe_shared_mlp_scales.py new file mode 100644 index 000000000000..702d62003101 --- /dev/null +++ b/tests/unittest/_torch/modules/test_megamoe_shared_mlp_scales.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""UE8M0 block scales that the fused shared FC12 kernel reads on SM100 and SM103.""" + +import pytest +import torch + +from tensorrt_llm._torch.modules.linear import Linear +from tensorrt_llm._torch.modules.megamoe_shared_mlp import _unpack_r128c4 +from tensorrt_llm._utils import get_sm_version, is_sm_100f +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo +from tensorrt_llm.quantization.utils.fp8_utils import transform_k128_scales_to_cutedsl_mxfp8_layout + +pytestmark = pytest.mark.skipif( + not is_sm_100f() or get_sm_version() == 107, + reason="SM107 converts these scales to the CuTe DSL MXFP8 layout instead", +) + + +def _block_scaled_weights(out_features: int, in_features: int) -> dict[str, torch.Tensor]: + weight = torch.randn(out_features, in_features) / in_features**0.5 + blocks = weight.reshape(out_features // 128, 128, in_features // 128, 128) + scale = blocks.abs().amax(dim=(1, 3)) / torch.finfo(torch.float8_e4m3fn).max + quantized = (blocks / scale[:, None, :, None]).reshape(out_features, in_features) + return {"weight": quantized.to(torch.float8_e4m3fn), "weight_scale": scale.float()} + + +def _dequantize(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + expanded = scale.repeat_interleave(128, dim=0).repeat_interleave(128, dim=1) + return weight.float() * expanded + + +@torch.inference_mode() +def test_linear_rounds_block_scales_to_ue8m0_when_asked(): + out_features, in_features = 256, 512 + weights = _block_scaled_weights(out_features, in_features) + linear = Linear( + in_features, + out_features, + bias=False, + dtype=torch.bfloat16, + quant_config=QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES), + use_cute_dsl_blockscaling_mm=True, + ) + linear.round_block_scales_to_ue8m0 = True + linear.load_weights([weights]) + linear.post_load_weights() + linear.cuda() + + scale = linear.weight_scale + assert scale.dtype == torch.float32 + assert tuple(scale.shape) == (out_features // 128, in_features // 128) + assert linear.weight.dtype == torch.float8_e4m3fn + torch.testing.assert_close(scale, torch.exp2(torch.round(torch.log2(scale))), rtol=0, atol=0) + # Rounding a block's scale up to a power of two rescales its FP8 values, so the + # dequantized weight changes by at most FP8 rounding. + torch.testing.assert_close( + _dequantize(linear.weight, scale), + _dequantize(weights["weight"].cuda(), weights["weight_scale"].cuda()), + rtol=0.07, + atol=1e-3, + ) + # The fused kernel reads these scales through the CuTe DSL MXFP8 layout. + packed = transform_k128_scales_to_cutedsl_mxfp8_layout(scale, mn=out_features, k=in_features) + unpacked = _unpack_r128c4(packed, out_features, in_features // 32) + expected = scale.repeat_interleave(128, dim=0).repeat_interleave(4, dim=1) + torch.testing.assert_close( + unpacked, expected.to(torch.float8_e8m0fnu).view(torch.uint8), rtol=0, atol=0 + ) + + +@torch.inference_mode() +def test_linear_keeps_block_scales_without_the_request(): + out_features, in_features = 256, 512 + weights = _block_scaled_weights(out_features, in_features) + linear = Linear( + in_features, + out_features, + bias=False, + dtype=torch.bfloat16, + quant_config=QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES), + use_cute_dsl_blockscaling_mm=True, + ) + linear.load_weights([weights]) + linear.post_load_weights() + linear.cuda() + + torch.testing.assert_close(linear.weight_scale.cpu(), weights["weight_scale"], rtol=0, atol=0) From 8730d8e1eeb17ea73b1122625aa3da4f2f7f3453 Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 20 Oct 2026 18:25:54 -0700 Subject: [PATCH 31/59] [None][perf] Default per-layer dynamic EPLB to four auxiliary SMs The fused HALO-Q + in-switch copy kernel runs one CTA per auxiliary SM, and the routed-input quantization and DeepSeek-V4's fused shared-expert kernel leave those SMs free. The default drops from 8 to 4 on Blackwell and Rubin alike, so those kernels keep four more SMs. The allowed range stays 1 to 8, and the serve-config schema snapshots carry the new default. Signed-off-by: jintaop --- examples/wide_ep/ep_load_balancer/README.md | 4 ++-- tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py | 2 +- tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py | 2 +- tensorrt_llm/llmapi/llm_args.py | 2 +- tensorrt_llm/schemas/trtllm-serve-config.schema.json | 2 +- tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json | 2 +- tests/unittest/llmapi/test_llm_args.py | 6 +++--- 7 files changed, 10 insertions(+), 10 deletions(-) diff --git a/examples/wide_ep/ep_load_balancer/README.md b/examples/wide_ep/ep_load_balancer/README.md index e8d8644334cf..36aa87b31230 100644 --- a/examples/wide_ep/ep_load_balancer/README.md +++ b/examples/wide_ep/ep_load_balancer/README.md @@ -294,7 +294,7 @@ moe_config: load_balancer: mode: per_layer_dynamic num_slots: 416 # 384 experts + 4 helper slots on each of 8 EP ranks - auxiliary_sms: 8 # optional; 8 is the default + auxiliary_sms: 4 # optional; 4 is the default helper_slot_mode: per_layer_resident # optional; the default EOF ``` @@ -309,7 +309,7 @@ Per-layer dynamic EPLB has the following constraints: * **Parallelism**: attention DP (`enable_attention_dp: true`) with pure expert parallelism: `moe_expert_parallel_size` equals the number of GPUs, with no MoE tensor parallelism and no pipeline parallelism. The EP size must be even and between 2 and 32, and the MoE smart router (`moe_cluster_size` greater than 1) is not supported. Unlike `standard` mode, an unsupported model, an EP size of 1, or the smart router raises an error instead of silently disabling the load balancer. * **CUDA graphs**: not supported. A non-null `cuda_graph_config` is replaced with `null` and a warning is logged, and any `prefill_cuda_graph_backend` other than `DISABLED` is rejected. * **Launch queues**: CUDA must be initialized with `CUDA_SCALE_LAUNCH_QUEUES=4x`. The LLM API sets it automatically, in the main process and in every worker, before CUDA is initialized. If the process has already initialized CUDA when the LLM is created, construction fails; export `CUDA_SCALE_LAUNCH_QUEUES=4x` before starting such a process. -* **`auxiliary_sms`**: the number of SMs of the kernel that runs the HALO-Q scheduler and then the copy, one CTA per SM. It defaults to 8, must be between 1 and 8, and must be lower than the GPU's SM count. Input quantization and, on the SM100 family, DeepSeek-V4's fused shared-expert kernel leave these SMs free, so planning and copying run alongside them. +* **`auxiliary_sms`**: the number of SMs of the kernel that runs the HALO-Q scheduler and then the copy, one CTA per SM. It defaults to 4, must be between 1 and 8, and must be lower than the GPU's SM count. Input quantization and, on the SM100 family, DeepSeek-V4's fused shared-expert kernel leave these SMs free, so planning and copying run alongside them. * **`helper_slot_mode`**: `per_layer_resident` (the default) gives every MoE layer its own helper slots and skips copying experts that they still hold from the previous forward pass. It costs one helper set per MoE layer, about 6.4 GiB per GPU for DeepSeek-V4-Pro with 3 helper slots per rank. `shared_banks` has all MoE layers alternate between two shared helper banks and copies every helper expert in every forward pass. * **DWDP**: not supported; `dwdp_config` must keep `dwdp_size` at 1. * **Fields**: only `mode`, `num_slots`, `auxiliary_sms` and `helper_slot_mode`. `num_slots` is required, must be greater than the expert count, and must be divisible by the EP size. `layer_updates_per_iter` and `initial_global_assignments` are `standard`-only and rejected in this mode. diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py b/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py index 69c642d64978..83467711a693 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py @@ -1145,7 +1145,7 @@ def maybe_create_moe_load_balancer( "Created per-layer dynamic MoE LoadBalancer, num_slots=%s, " "auxiliary_sms=%s...", config.num_slots, - getattr(config, "auxiliary_sms", 8), + getattr(config, "auxiliary_sms", 4), ) elif mode == "standard": if config.layer_updates_per_iter > 0: diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py index 56604ce5f690..7f3f53078c45 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py @@ -386,7 +386,7 @@ def __init__( self.mapping = mapping self.num_slots = num_slots self.num_local_slots = num_slots // ep_size - self.auxiliary_sms = int(getattr(config, "auxiliary_sms", 8)) + self.auxiliary_sms = int(getattr(config, "auxiliary_sms", 4)) if self.auxiliary_sms <= 0: raise ValueError("auxiliary_sms must be positive") helper_slot_mode = getattr(config, "helper_slot_mode", "per_layer_resident") diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 011743a8b92b..13a6dd998346 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1956,7 +1956,7 @@ class MoePerLayerDynamicLoadBalancerConfig(_MoeLoadBalancerConfigBase): "Total number of expert slots across all EP ranks. Must be divisible by `ep_size` and greater than the number of experts. Each rank keeps its `num_experts / ep_size` resident experts and adds `(num_slots - num_experts) / ep_size` helper slots. Unlike standard EPLB, the helper slots are not static replicas: every forward pass refills them from that layer's routing; see `helper_slot_mode` for their storage.", status="prototype") auxiliary_sms: int = Field( - default=8, + default=4, ge=1, le=8, description= diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index 2d2fca14a461..c71d0d85eef0 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -4133,7 +4133,7 @@ "description": "Configure per-layer dynamic EPLB with Halo-Q scheduling.", "properties": { "auxiliary_sms": { - "default": 8, + "default": 4, "description": "Number of SMs that the routed-input quantization and shared-expert FC12 kernels leave free for the kernel that runs the HALO-Q scheduler and then the in-switch copy of its plan on the copy stream, one CTA per SM. Must be smaller than the device SM count.", "maximum": 8, "minimum": 1, diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index 7ca5daa75c17..fe9ed51a6e11 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -5697,7 +5697,7 @@ "description": "Configure per-layer dynamic EPLB with Halo-Q scheduling.", "properties": { "auxiliary_sms": { - "default": 8, + "default": 4, "description": "Number of SMs that the routed-input quantization and shared-expert FC12 kernels leave free for the kernel that runs the HALO-Q scheduler and then the in-switch copy of its plan on the copy stream, one CTA per SM. Must be smaller than the device SM count.", "maximum": 8, "minimum": 1, diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 5cecd7f6eee6..48e2253a21ed 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -274,7 +274,7 @@ def test_resolves_auto_and_accepts_compatible_backend( assert isinstance(config.load_balancer, MoePerLayerDynamicLoadBalancerConfig) assert config.backend == "MEGAMOE_CUTEDSL" - assert config.load_balancer.auxiliary_sms == 8 + assert config.load_balancer.auxiliary_sms == 4 assert config.load_balancer.helper_slot_mode == "per_layer_resident" @pytest.mark.parametrize( @@ -327,8 +327,8 @@ def test_per_layer_dynamic_positive_fields(self, field: str) -> None: with pytest.raises(ValidationError, match=field): MoeConfig(load_balancer=values) - @pytest.mark.parametrize("auxiliary_sms", [1, 4]) - def test_per_layer_dynamic_accepts_smaller_auxiliary_sm_budget( + @pytest.mark.parametrize("auxiliary_sms", [1, 8]) + def test_per_layer_dynamic_accepts_other_auxiliary_sm_budgets( self, auxiliary_sms: int) -> None: values = self._per_layer_dynamic() values["auxiliary_sms"] = auxiliary_sms From d1a4325ca355f0d2974c236850c27900c0505fad Mon Sep 17 00:00:00 2001 From: jintaop Date: Wed, 21 Oct 2026 02:52:23 -0700 Subject: [PATCH 32/59] [None][chore] Remove the standalone HALO-Q kernel; keep only the fused launch The fused HALO-Q + in-switch copy kernel is the only scheduler launch, so the standalone scheduler path is dead code. Removed: - halo_q_scheduler_kernel, invokeMoeRebalanceHaloQ, and the moe_rebalance_halo_q op with its fake registration. - The scheduler body's plan-channel publication and PDL trigger, which only that kernel used, and their arguments (plan channel pointer, plan ABI version and words, route features, enable_pdl) in MoeRebalanceHaloQParams and moe_rebalance_tma_submit_fused. - CudaPhysicalSlotScheduler's stage/launch/submit/run path and its plan channel binding. bind_gpu_direct and fused_submit_args remain. - The scheduler group's one-time standalone launch at construction. A barrier still keeps every rank's exchange buffer mapped before any rank can launch the fused kernel, and fused launches during warmup and autotuning already use the warmup spin bound. Signed-off-by: jintaop --- .../dynamicEplb/moeRebalanceHaloQ.cu | 124 ------ .../dynamicEplb/moeRebalanceHaloQ.h | 11 +- .../dynamicEplb/moeRebalanceHaloQKernel.cuh | 51 +-- .../dynamicEplb/moeRebalancePlaceholder.cpp | 5 - .../dynamicEplb/moeRebalanceTma.cu | 12 +- .../loadBalance/dynamicEplb/moeRebalanceTma.h | 2 +- .../moe/loadBalance/moeRebalanceHaloQOp.cpp | 56 +-- .../moe/loadBalance/moeRebalanceHaloQOp.h | 10 +- .../moe/loadBalance/moeRebalanceTmaOp.cpp | 15 +- .../_torch/custom_ops/cpp_custom_ops.py | 31 -- .../mega_moe/rebalance_slot_scheduler.py | 13 +- .../cuda_scheduler/runtime.py | 370 +++--------------- .../_torch/moe/test_dynamic_eplb_contracts.py | 69 +--- 13 files changed, 105 insertions(+), 664 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.cu b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.cu index f48c438602c6..59d440ba29f3 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.cu +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.cu @@ -15,77 +15,16 @@ */ #include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h" #include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh" -#include - #include -#include -#include -#include -#include TRTLLM_NAMESPACE_BEGIN namespace kernels { -namespace -{ - -template -__global__ __launch_bounds__(512, 1) void halo_q_scheduler_kernel(int const* routes, int* out_slots, int* out_ids, - int* out_levels, int* out_owners, unsigned long long const* peer_bases, int* status, int* partial, int* route_aux, - int* grid_sync, int* plan, int* route_prefix, int* plan_channel, int runtime_ep, int runtime_experts, - int runtime_helpers, int local_rank, int route_capacity, int ctas, int enable_pdl, unsigned long long spin_cycles, - int plan_abi_version, int plan_channel_words, int route_features, int capacity, int route_count) -{ - halo_q_scheduler_body(routes, out_slots, out_ids, out_levels, out_owners, - peer_bases, status, partial, route_aux, grid_sync, plan, route_prefix, plan_channel, runtime_ep, - runtime_experts, runtime_helpers, local_rank, route_capacity, ctas, enable_pdl, spin_cycles, plan_abi_version, - plan_channel_words, route_features, capacity, route_count); -} - -bool validPlanChannelLaunch(std::uintptr_t pointer, int abiVersion, int words, int features, int ep, int helpers) -{ - if (pointer == 0) - { - return abiVersion == 0 && words == 0 && features == 0; - } - if (pointer % 16 != 0 || helpers <= 0) - { - return false; - } - auto const stride = (static_cast(helpers) + 3) & ~std::int64_t{3}; - auto const expected = 4 + 7 * stride; - return abiVersion == 6 && ep >= 2 && ep <= kMaxEp && ep % 2 == 0 && (features == 0 || features == 1) - && expected <= INT_MAX && words == expected; -} - -cudaError_t ensureDynamicSharedMemory(void const* kernel, int sharedBytes) -{ - int device = -1; - auto error = cudaGetDevice(&device); - if (error != cudaSuccess) - return error; - - using KernelCache = std::map>; - static std::mutex mutex; - static std::map configuredBytes; - std::lock_guard lock(mutex); - auto& current = configuredBytes[device][kernel]; - if (current >= sharedBytes) - return cudaSuccess; - error = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, sharedBytes); - if (error == cudaSuccess) - current = sharedBytes; - return error; -} - -} // namespace - void checkMoeRebalanceHaloQParams(MoeRebalanceHaloQParams const& params) { TLLM_CHECK_WITH_INFO((params.routes != nullptr || params.validRouteCount == 0) && params.outSlots != nullptr @@ -100,69 +39,6 @@ void checkMoeRebalanceHaloQParams(MoeRebalanceHaloQParams const& params) && params.routeCapacity <= INT_MAX / params.ep && params.routeCapacity <= INT_MAX - params.ctas && params.ctas > 0 && params.ctas <= 128 && params.threads == 512, "Invalid HALO-Q scheduler launch parameters"); - TLLM_CHECK_WITH_INFO( - validPlanChannelLaunch(reinterpret_cast(params.planChannel), params.planAbiVersion, - params.planChannelWords, params.routeFeatures, params.ep, params.helpers), - "Invalid HALO-Q PlanChannel launch ABI"); -} - -void invokeMoeRebalanceHaloQ(MoeRebalanceHaloQParams const& params, cudaStream_t stream) -{ - checkMoeRebalanceHaloQParams(params); - - int capacity = broadcast_capacity(params.ep, params.experts, params.helpers); - std::size_t const sharedBytes - = static_cast(halo_q_shared_ints(params.ep, params.experts, capacity, params.threads)) - * sizeof(int); - - void const* kernel = reinterpret_cast(&halo_q_scheduler_kernel<0, 0, 0>); - if (params.experts == 384 && params.helpers == 4) - { - if (params.ep == 8) - { - kernel = reinterpret_cast(&halo_q_scheduler_kernel<8, 384, 4>); - } - else if (params.ep == 4) - { - kernel = reinterpret_cast(&halo_q_scheduler_kernel<4, 384, 4>); - } - } - if (sharedBytes > 48 * 1024) - { - TLLM_CUDA_CHECK(ensureDynamicSharedMemory(kernel, static_cast(sharedBytes))); - } - - auto routes = params.routes; - auto outSlots = params.outSlots; - auto outIds = params.outIds; - auto outLevels = params.outLevels; - auto outOwners = params.outOwners; - auto peerBases = params.peerBases; - auto status = params.status; - auto partial = params.partial; - auto routeAux = params.routeAux; - auto gridSync = params.gridSync; - auto planWorkspace = params.planWorkspace; - auto routePrefix = params.routePrefix; - auto planChannel = params.planChannel; - int ep = params.ep; - int experts = params.experts; - int helpers = params.helpers; - int localRank = params.localRank; - int routeCapacity = params.routeCapacity; - int ctas = params.ctas; - int enablePdl = params.enablePdl ? 1 : 0; - std::uint64_t spinCycles = params.spinCycles; - int planAbiVersion = params.planAbiVersion; - int planChannelWords = params.planChannelWords; - int routeFeatures = params.routeFeatures; - int validRouteCount = params.validRouteCount; - void* arguments[] - = {&routes, &outSlots, &outIds, &outLevels, &outOwners, &peerBases, &status, &partial, &routeAux, &gridSync, - &planWorkspace, &routePrefix, &planChannel, &ep, &experts, &helpers, &localRank, &routeCapacity, &ctas, - &enablePdl, &spinCycles, &planAbiVersion, &planChannelWords, &routeFeatures, &capacity, &validRouteCount}; - TLLM_CUDA_CHECK( - cudaLaunchKernel(kernel, dim3(ctas, 1, 1), dim3(params.threads, 1, 1), arguments, sharedBytes, stream)); } } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h index 5d0d10a62bb6..01901208d624 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h @@ -19,8 +19,6 @@ #include "tensorrt_llm/common/config.h" -#include - #include TRTLLM_NAMESPACE_BEGIN @@ -28,6 +26,8 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { +// The HALO-Q scheduler phase of the fused scheduler + in-switch copy kernel +// (megamoe_tma_copy_submit_fused). struct MoeRebalanceHaloQParams { int const* routes; @@ -42,7 +42,6 @@ struct MoeRebalanceHaloQParams int* gridSync; int* planWorkspace; int* routePrefix; - int* planChannel; int ep; int experts; int helpers; @@ -50,19 +49,13 @@ struct MoeRebalanceHaloQParams int routeCapacity; int ctas; int threads; - bool enablePdl; std::uint64_t spinCycles; - int planAbiVersion; - int planChannelWords; - int routeFeatures; int validRouteCount; }; // Throws unless params describe a valid scheduler launch. void checkMoeRebalanceHaloQParams(MoeRebalanceHaloQParams const& params); -void invokeMoeRebalanceHaloQ(MoeRebalanceHaloQParams const& params, cudaStream_t stream); - } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh index 970fac3420d4..d71491dcfc13 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh @@ -14,8 +14,8 @@ * limitations under the License. */ -/* HALO-Q scheduler device code. The standalone scheduler kernel and the fused - * scheduler + in-switch copy kernel both run halo_q_scheduler_body. */ +/* HALO-Q scheduler device code, the first phase of the fused scheduler + + * in-switch copy kernel. */ #ifndef TRTLLM_KERNELS_MOE_REBALANCE_HALO_Q_KERNEL_CUH #define TRTLLM_KERNELS_MOE_REBALANCE_HALO_Q_KERNEL_CUH @@ -54,15 +54,6 @@ enum PlanField kPlanHelper = 5, }; -// Adapt placement storage to the copy layer without exposing scheduler -// internals. -__device__ __forceinline__ sami_copy::PlacementView copy_placement_view(int const* plan) -{ - return {plan[1], kMaxBroadcasts, plan + 2 + kPlanExpert * kMaxBroadcasts, plan + 2 + kPlanLevel * kMaxBroadcasts, - plan + 2 + kPlanGroupBegin * kMaxBroadcasts, plan + 2 + kPlanGroupSize * kMaxBroadcasts, - plan + 2 + kPlanOwner * kMaxBroadcasts, plan + 2 + kPlanHelper * kMaxBroadcasts}; -} - __device__ __forceinline__ int imin(int a, int b) { return a < b ? a : b; @@ -1774,25 +1765,16 @@ __device__ void materialize_worker_routes(int const* routes, int* output, int co /* One scheduler generation on a grid of ctas resident CTAs of 512 threads: * CTA 0 plans and publishes the helper placement, the other CTAs write routes. - * Callers launch it as halo_q_scheduler_kernel or as the first phase of the - * fused scheduler + in-switch copy kernel. */ + * The fused scheduler + in-switch copy kernel runs it as its first phase. */ template __device__ __forceinline__ void halo_q_scheduler_body(int const* routes, int* out_slots, int* out_ids, int* out_levels, int* out_owners, unsigned long long const* peer_bases, int* status, int* partial, int* route_aux, int* grid_sync, - int* plan, int* route_prefix, int* plan_channel, int runtime_ep, int runtime_experts, int runtime_helpers, - int local_rank, int route_capacity, int ctas, int enable_pdl, unsigned long long spin_cycles, int plan_abi_version, - int plan_channel_words, int route_features, int capacity, int route_count) + int* plan, int* route_prefix, int runtime_ep, int runtime_experts, int runtime_helpers, int local_rank, + int route_capacity, int ctas, unsigned long long spin_cycles, int capacity, int route_count) { int const ep = FixedEp > 0 ? FixedEp : runtime_ep; int const experts = FixedExperts > 0 ? FixedExperts : runtime_experts; int const helpers = FixedHelpers > 0 ? FixedHelpers : runtime_helpers; - (void) plan_channel_words; // Validated by the host launch wrapper. - /* The scheduler grid is deliberately residency-safe (one CTA per SM and - * CTA count <= SM count), so signal as soon as every CTA starts. The - * independent same-stream quantizer may then occupy all unused SMs while - * these CTAs retain their own SMs for the complete scheduler lifetime. */ - if (enable_pdl && threadIdx.x == 0) - cudaTriggerProgrammaticLaunchCompletion(); // Only the caller's valid prefix is readable. Keep the public output at // its configured capacity, clearing stale routes inside this same kernel. // Empty ranks still execute every collective; no route load occurs at N=0. @@ -1911,31 +1893,16 @@ __device__ __forceinline__ void halo_q_scheduler_body(int const* routes, int* ou } if (warp == 0 || warp == 8) scheduler_tail_barrier(); - if (warp == 0) - { - /* Release the worker CTAs here rather than after the channel - * publish. They gate on grid_sync+2 and consume route_prefix, - * which is complete above; plan_channel is consumed only by the - * copy stack, which acquires on the channel's own epoch word - * written with store_release_system below. Publishing after - * the release lets the workers coloring pass -- the work that - * actually determines when the kernel ends -- overlap the whole - * publish instead of waiting behind it. */ - if (threadIdx.x == 0) - store_release_device(grid_sync + 2, epoch); - sami_copy::publish_plan_channel_warp(plan_channel, out_ids, out_levels, out_owners, - copy_placement_view(plan), helpers, ep, local_rank, experts, plan_abi_version, route_features); - __syncwarp(); - } + /* The worker CTAs gate on grid_sync+2 and consume route_prefix, + * which is complete above. */ + if (threadIdx.x == 0) + store_release_device(grid_sync + 2, epoch); return; } if (threadIdx.x == 0) store_release_device(plan + 0, epoch); __syncthreads(); build_route_prefix(counts, quota, route_prefix, plan, broadcasts, ep, experts, local_rank, blockDim.x); - __syncthreads(); - sami_copy::publish_plan_channel_warp(plan_channel, out_ids, out_levels, out_owners, copy_placement_view(plan), - helpers, ep, local_rank, experts, plan_abi_version, route_features); /* The block barrier orders route-prefix/helper writes before release. */ __syncthreads(); } diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlaceholder.cpp b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlaceholder.cpp index 259820957bf9..929a18cce151 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlaceholder.cpp +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlaceholder.cpp @@ -33,11 +33,6 @@ void checkMoeRebalanceHaloQParams(MoeRebalanceHaloQParams const&) TLLM_THROW("Per-layer dynamic EPLB HALO-Q requires a requested SM100, SM103, or SM107 target"); } -void invokeMoeRebalanceHaloQ(MoeRebalanceHaloQParams const&, cudaStream_t) -{ - TLLM_THROW("Per-layer dynamic EPLB HALO-Q requires a requested SM100, SM103, or SM107 target"); -} - } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu index d2d5fe82c6d6..adec20d46c8a 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu @@ -752,8 +752,8 @@ __global__ void __launch_bounds__(512, 1) halo_q_tma_fused_kernel(FusedScheduler tensorrt_llm::kernels::halo_q_scheduler_body(scheduler.routes, scheduler.out_slots, scheduler.out_ids, scheduler.out_levels, scheduler.out_owners, scheduler.peer_bases, scheduler.status, scheduler.partial, scheduler.route_aux, scheduler.grid_sync, scheduler.plan, - scheduler.route_prefix, nullptr, scheduler.ep, scheduler.experts, scheduler.helpers, scheduler.local_rank, - scheduler.route_capacity, static_cast(gridDim.x), 0, scheduler.spin_cycles, 0, 0, 0, scheduler.capacity, + scheduler.route_prefix, scheduler.ep, scheduler.experts, scheduler.helpers, scheduler.local_rank, + scheduler.route_capacity, static_cast(gridDim.x), scheduler.spin_cycles, scheduler.capacity, scheduler.route_count); // Orders every thread's scheduler writes before thread 0's release. __syncthreads(); @@ -1246,10 +1246,10 @@ int megamoe_tma_copy_submit_fused(MegamoeTmaCopyState* state, auto const& d = state->direct; // The copy reads the bound scheduler outputs, so the scheduler must write exactly those. if (!state->direct_bound || scheduler.ctas != c.sms || scheduler.threads != kFusedThreads - || scheduler.planChannel != nullptr || scheduler.outIds != d.ids || scheduler.outLevels != d.levels - || scheduler.outOwners != d.owners || scheduler.planWorkspace != d.workspace || scheduler.status != d.status - || scheduler.gridSync != d.grid_sync || scheduler.ep != d.config.world - || scheduler.experts != d.config.global_experts || scheduler.helpers != d.config.helper_count) + || scheduler.outIds != d.ids || scheduler.outLevels != d.levels || scheduler.outOwners != d.owners + || scheduler.planWorkspace != d.workspace || scheduler.status != d.status || scheduler.gridSync != d.grid_sync + || scheduler.ep != d.config.world || scheduler.experts != d.config.global_experts + || scheduler.helpers != d.config.helper_count) return cudaErrorInvalidValue; int device = -1; cudaError_t error = cudaGetDevice(&device); diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h index 582e6b3c6594..39668c395f18 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h @@ -206,7 +206,7 @@ extern "C" /* Launches HALO-Q and the GPU-direct copy of its plan as one kernel of config.sms * CTAs. scheduler must pass checkMoeRebalanceHaloQParams, use config.sms CTAs and - * no plan channel, and write the outputs bound by megamoe_tma_copy_bind_gpu_direct. + * write the outputs bound by megamoe_tma_copy_bind_gpu_direct. * route_epoch is a caller-owned uint64 that only fused launches change: each CTA * adds one with a system-scope release once its routes are written, and * route_target is its value after this launch. A cuStreamWaitValue64 (GEQ diff --git a/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.cpp b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.cpp index 33b8cfa0855f..334e72cf9332 100644 --- a/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.cpp +++ b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.cpp @@ -16,8 +16,6 @@ #include "tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.h" -#include -#include #include #include @@ -54,11 +52,9 @@ int checkedInt(std::int64_t value, char const* name) kernels::MoeRebalanceHaloQParams makeMoeRebalanceHaloQParams(torch::Tensor const& routes, torch::Tensor& outSlots, torch::Tensor& outIds, torch::Tensor& outLevels, torch::Tensor& outOwners, torch::Tensor const& peerBases, torch::Tensor& status, torch::Tensor& partial, torch::Tensor& routeAux, torch::Tensor& gridSync, - torch::Tensor& planWorkspace, torch::Tensor& routePrefix, std::int64_t planChannelPtr, std::int64_t epValue, - std::int64_t expertsValue, std::int64_t helpersValue, std::int64_t localRankValue, std::int64_t routeCapacityValue, - std::int64_t ctasValue, std::int64_t threadsValue, bool enablePdl, std::int64_t spinCycles, - std::int64_t planAbiVersionValue, std::int64_t planChannelWordsValue, std::int64_t routeFeaturesValue, - std::int64_t validRouteCountValue) + torch::Tensor& planWorkspace, torch::Tensor& routePrefix, std::int64_t epValue, std::int64_t expertsValue, + std::int64_t helpersValue, std::int64_t localRankValue, std::int64_t routeCapacityValue, std::int64_t ctasValue, + std::int64_t threadsValue, std::int64_t spinCycles, std::int64_t validRouteCountValue) { int const ep = checkedInt(epValue, "ep"); int const experts = checkedInt(expertsValue, "experts"); @@ -67,9 +63,6 @@ kernels::MoeRebalanceHaloQParams makeMoeRebalanceHaloQParams(torch::Tensor const int const routeCapacity = checkedInt(routeCapacityValue, "route_capacity"); int const ctas = checkedInt(ctasValue, "ctas"); int const threads = checkedInt(threadsValue, "threads"); - int const planAbiVersion = checkedInt(planAbiVersionValue, "plan_abi_version"); - int const planChannelWords = checkedInt(planChannelWordsValue, "plan_channel_words"); - int const routeFeatures = checkedInt(routeFeaturesValue, "route_features"); int const validRouteCount = checkedInt(validRouteCountValue, "valid_route_count"); TORCH_CHECK(ep >= 2 && ep <= 32 && ep % 2 == 0, "ep must be even and in [2, 32]"); TORCH_CHECK(experts > 0 && experts <= 384 && experts % ep == 0, @@ -82,7 +75,6 @@ kernels::MoeRebalanceHaloQParams makeMoeRebalanceHaloQParams(torch::Tensor const TORCH_CHECK(ctas > 0 && ctas <= 128 && routeCapacity <= INT_MAX - ctas, "invalid ctas"); TORCH_CHECK(threads == 512, "threads must be 512"); TORCH_CHECK(spinCycles >= 0, "spin_cycles cannot be negative"); - TORCH_CHECK(planChannelPtr >= 0, "plan_channel_ptr cannot be negative"); c10::Device const device = routes.device(); checkCudaTensor(routes, device, torch::kInt32, validRouteCount, "routes"); @@ -117,7 +109,6 @@ kernels::MoeRebalanceHaloQParams makeMoeRebalanceHaloQParams(torch::Tensor const gridSync.data_ptr(), planWorkspace.data_ptr(), routePrefix.data_ptr(), - reinterpret_cast(static_cast(planChannelPtr)), ep, experts, helpers, @@ -125,52 +116,11 @@ kernels::MoeRebalanceHaloQParams makeMoeRebalanceHaloQParams(torch::Tensor const routeCapacity, ctas, threads, - enablePdl, static_cast(spinCycles), - planAbiVersion, - planChannelWords, - routeFeatures, validRouteCount, }; } -void moeRebalanceHaloQ(torch::Tensor const& routes, torch::Tensor& outSlots, torch::Tensor& outIds, - torch::Tensor& outLevels, torch::Tensor& outOwners, torch::Tensor const& peerBases, torch::Tensor& status, - torch::Tensor& partial, torch::Tensor& routeAux, torch::Tensor& gridSync, torch::Tensor& planWorkspace, - torch::Tensor& routePrefix, std::int64_t planChannelPtr, std::int64_t ep, std::int64_t experts, - std::int64_t helpers, std::int64_t localRank, std::int64_t routeCapacity, std::int64_t ctas, std::int64_t threads, - bool enablePdl, std::int64_t spinCycles, std::int64_t planAbiVersion, std::int64_t planChannelWords, - std::int64_t routeFeatures, std::int64_t validRouteCount) -{ - auto const params = makeMoeRebalanceHaloQParams(routes, outSlots, outIds, outLevels, outOwners, peerBases, status, - partial, routeAux, gridSync, planWorkspace, routePrefix, planChannelPtr, ep, experts, helpers, localRank, - routeCapacity, ctas, threads, enablePdl, spinCycles, planAbiVersion, planChannelWords, routeFeatures, - validRouteCount); - c10::cuda::CUDAGuard const deviceGuard(routes.device()); - kernels::invokeMoeRebalanceHaloQ(params, at::cuda::getCurrentCUDAStream(routes.get_device())); -} - } // namespace torch_ext TRTLLM_NAMESPACE_END - -TORCH_LIBRARY_FRAGMENT(trtllm, m) -{ - m.def( - "moe_rebalance_halo_q(Tensor routes, Tensor(a!) out_slots, Tensor(b!) " - "out_ids, Tensor(c!) out_levels, " - "Tensor(d!) out_owners, Tensor peer_bases, Tensor(e!) status, " - "Tensor(f!) partial, Tensor(g!) route_aux, " - "Tensor(h!) grid_sync, Tensor(i!) plan_workspace, Tensor(j!) " - "route_prefix, int plan_channel_ptr, int ep, " - "int experts, int helpers, int local_rank, int route_capacity, int " - "ctas, int threads, bool enable_pdl, " - "int spin_cycles, int plan_abi_version, int plan_channel_words, int " - "route_features, " - "int valid_route_count) -> ()"); -} - -TORCH_LIBRARY_IMPL(trtllm, CUDA, m) -{ - m.impl("moe_rebalance_halo_q", &tensorrt_llm::torch_ext::moeRebalanceHaloQ); -} diff --git a/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.h b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.h index 71213603b6e9..d3bef11930b2 100644 --- a/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.h +++ b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.h @@ -27,15 +27,13 @@ TRTLLM_NAMESPACE_BEGIN namespace torch_ext { -// Validates the moe_rebalance_halo_q arguments and returns its launch parameters. -// The fused scheduler + in-switch copy op takes the same scheduler arguments. +// Validates the HALO-Q scheduler arguments of moe_rebalance_tma_submit_fused and returns its launch parameters. kernels::MoeRebalanceHaloQParams makeMoeRebalanceHaloQParams(torch::Tensor const& routes, torch::Tensor& outSlots, torch::Tensor& outIds, torch::Tensor& outLevels, torch::Tensor& outOwners, torch::Tensor const& peerBases, torch::Tensor& status, torch::Tensor& partial, torch::Tensor& routeAux, torch::Tensor& gridSync, - torch::Tensor& planWorkspace, torch::Tensor& routePrefix, std::int64_t planChannelPtr, std::int64_t ep, - std::int64_t experts, std::int64_t helpers, std::int64_t localRank, std::int64_t routeCapacity, std::int64_t ctas, - std::int64_t threads, bool enablePdl, std::int64_t spinCycles, std::int64_t planAbiVersion, - std::int64_t planChannelWords, std::int64_t routeFeatures, std::int64_t validRouteCount); + torch::Tensor& planWorkspace, torch::Tensor& routePrefix, std::int64_t ep, std::int64_t experts, + std::int64_t helpers, std::int64_t localRank, std::int64_t routeCapacity, std::int64_t ctas, std::int64_t threads, + std::int64_t spinCycles, std::int64_t validRouteCount); } // namespace torch_ext diff --git a/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp index c81984897757..39f34ca6bc74 100644 --- a/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp +++ b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp @@ -383,15 +383,13 @@ void moeRebalanceTmaSubmitFused(MoeRebalanceTmaStatePtr const& state, int64_t fl int64_t routeTarget, torch::Tensor const& routes, torch::Tensor& outSlots, torch::Tensor& outIds, torch::Tensor& outLevels, torch::Tensor& outOwners, torch::Tensor const& peerBases, torch::Tensor& status, torch::Tensor& partial, torch::Tensor& routeAux, torch::Tensor& gridSync, torch::Tensor& planWorkspace, - torch::Tensor& routePrefix, int64_t planChannelPtr, int64_t ep, int64_t experts, int64_t helpers, int64_t localRank, - int64_t routeCapacity, int64_t ctas, int64_t threads, bool enablePdl, int64_t spinCycles, int64_t planAbiVersion, - int64_t planChannelWords, int64_t routeFeatures, int64_t validRouteCount) + torch::Tensor& routePrefix, int64_t ep, int64_t experts, int64_t helpers, int64_t localRank, int64_t routeCapacity, + int64_t ctas, int64_t threads, int64_t spinCycles, int64_t validRouteCount) { TORCH_CHECK(state, "moe rebalance TMA state is null"); auto const scheduler = makeMoeRebalanceHaloQParams(routes, outSlots, outIds, outLevels, outOwners, peerBases, - status, partial, routeAux, gridSync, planWorkspace, routePrefix, planChannelPtr, ep, experts, helpers, - localRank, routeCapacity, ctas, threads, enablePdl, spinCycles, planAbiVersion, planChannelWords, routeFeatures, - validRouteCount); + status, partial, routeAux, gridSync, planWorkspace, routePrefix, ep, experts, helpers, localRank, routeCapacity, + ctas, threads, spinCycles, validRouteCount); state->submitFused(flagMc, scheduler, routeEpoch, routeTarget); } @@ -445,9 +443,8 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state, int flag_mc, Tensor(a!) route_epoch, " "int route_target, Tensor routes, Tensor(b!) out_slots, Tensor(c!) out_ids, Tensor(d!) out_levels, " "Tensor(e!) out_owners, Tensor peer_bases, Tensor(f!) status, Tensor(g!) partial, Tensor(h!) route_aux, " - "Tensor(i!) grid_sync, Tensor(j!) plan_workspace, Tensor(k!) route_prefix, int plan_channel_ptr, int ep, " - "int experts, int helpers, int local_rank, int route_capacity, int ctas, int threads, bool enable_pdl, " - "int spin_cycles, int plan_abi_version, int plan_channel_words, int route_features, " + "Tensor(i!) grid_sync, Tensor(j!) plan_workspace, Tensor(k!) route_prefix, int ep, int experts, " + "int helpers, int local_rank, int route_capacity, int ctas, int threads, int spin_cycles, " "int valid_route_count) -> ()"); m.def( "moe_rebalance_tma_current_gen(" diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 394c5625d371..8657e26ee98d 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -72,37 +72,6 @@ def copy_kv_slots_( ) -> None: return None - @torch.library.register_fake("trtllm::moe_rebalance_halo_q") - def _( - routes: torch.Tensor, - out_slots: torch.Tensor, - out_ids: torch.Tensor, - out_levels: torch.Tensor, - out_owners: torch.Tensor, - peer_bases: torch.Tensor, - status: torch.Tensor, - partial: torch.Tensor, - route_aux: torch.Tensor, - grid_sync: torch.Tensor, - plan_workspace: torch.Tensor, - route_prefix: torch.Tensor, - plan_channel_ptr: int, - ep: int, - experts: int, - helpers: int, - local_rank: int, - route_capacity: int, - ctas: int, - threads: int, - enable_pdl: bool, - spin_cycles: int, - plan_abi_version: int, - plan_channel_words: int, - route_features: int, - valid_route_count: int, - ) -> None: - return None - @torch.library.register_fake("trtllm::allreduce") def allreduce( input: torch.Tensor, diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py index 693f3b042e71..2a25de9e95d5 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py @@ -231,7 +231,7 @@ def __init__( self._route_epoch_ptr = int(self._route_epoch.data_ptr()) self._route_target = 0 - # Every peer address must be connected before either warmup or submit; + # Every peer address must be connected before the first fused launch; # HALO-Q performs an all-rank exchange on the COPY stream. self.symmetric = FabricSymmetricBuffer( 4 * symmetric_buffer_ints(ep_size, logical_expert_count), self.device @@ -249,13 +249,10 @@ def __init__( ] ) - # Warm all ranks before binding the plan channel: a bound launch would - # leave a pending publication that must be consumed by the broadcaster. - # The scheduler input is already initialized and peer allocations synchronized. - # Loading the kernel can still skew the ranks, so this launch always - # waits for its peers as long as warmup does. - self.scheduler.launch(self.scheduler_stream, spin_cycles=spin_bound.warmup_cycles) - torch.cuda.synchronize(self.device) + # connect_peer_bases() synchronized this device, so its zeroed exchange + # buffer and peer mappings are complete. The barrier extends that to every + # rank before any rank binds its copy and can launch the fused kernel, + # whose HALO-Q exchange writes into its peers' buffers. ep_comm.barrier() # The provider must return the existing arena with identical bundle geometry. diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py index 0bd68b7fb0fe..10f0fb48ef84 100644 --- a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Host runtime for the pure-CUDA HALO-Q scheduler.""" +"""Host runtime for the HALO-Q scheduler, launched with its GPU-direct TMA copy as one kernel.""" from __future__ import annotations @@ -38,29 +38,6 @@ def _check_spin_cycles(spin_cycles: int) -> None: if type(spin_cycles) is not int or spin_cycles <= 0: raise ValueError("spin_cycles must be a positive exact int") -def _plan_channel_word_count(*, abi_version: int, channel_words: int | None, - route_features: int, ep: int, helpers: int) -> int: - """Validate an explicitly negotiated private publication allocation.""" - if type(abi_version) is not int or type(route_features) is not int: - raise ValueError("plan channel ABI and features must be exact integers") - if abi_version != 6: - raise ValueError("CUDA scheduler requires PlanChannel ABI6") - if route_features not in (0, 1): - raise ValueError("plan channel route features must be 0 or 1") - if type(ep) is not int or not MIN_EP <= ep <= MAX_EP or ep % 2: - raise ValueError("plan channel requires even EP in [2,32]") - if type(helpers) is not int or helpers <= 0: - raise ValueError("helper count must be a positive exact int") - stride = (helpers + 3) // 4 * 4 - expected = 4 + 7 * stride - if expected > 2**31 - 1: - raise ValueError("plan channel word count does not fit its int32 ABI") - if channel_words is None: - channel_words = expected - if type(channel_words) is not int or channel_words != expected: - raise ValueError("plan channel allocation does not match its explicit ABI") - return channel_words - def recommend_cuda_scheduler_ctas( ep_size: int, @@ -111,7 +88,6 @@ class CudaSchedulerConfig: local_rank: int = 0 threads: int = 512 ctas: int | None = None - enable_pdl: bool = True def __post_init__(self) -> None: if self.ctas is None: @@ -189,8 +165,6 @@ def validate(self) -> None: raise ValueError("the pure-CUDA scheduler currently requires threads=512") if not 1 <= self.ctas <= MAX_CTAS: raise ValueError(f"ctas must be in [1,{MAX_CTAS}]") - if type(self.enable_pdl) is not bool: - raise ValueError("enable_pdl must be bool") class CudaSchedulerOutputs: @@ -260,7 +234,12 @@ def _stream_handle(stream: object) -> int: class CudaPhysicalSlotScheduler: - """Preallocated pure-CUDA scheduler with no CuTe JIT executor.""" + """Preallocated HALO-Q outputs and workspace for the fused scheduler + copy kernel. + + The GPU-direct copy endpoint binds it once (``bind_gpu_direct``) and + launches it with the copy (``moe_rebalance_tma_submit_fused``), taking the + scheduler's arguments from ``fused_submit_args``. + """ def __init__(self, config: CudaSchedulerConfig, device: str | torch.device = "cuda"): config.validate() @@ -277,16 +256,9 @@ def __init__(self, config: CudaSchedulerConfig, device: str | torch.device = "cu self.device = scheduler_device self.outputs = CudaSchedulerOutputs(config, scheduler_device) self.out = self.outputs - self.logical_expert_ids = torch.full( - (config.max_tokens_per_rank, config.topk), - -1, - dtype=torch.int32, - device=scheduler_device, - ) self.status = torch.zeros(8, dtype=torch.int32, device=scheduler_device) - # Everything the kernel takes between `routes` and `stream` is fixed for - # the endpoint's lifetime. Cache it so a submission does not rebuild 13 - # data_ptr() calls and ~11 config lookups every generation. + # Every scheduler argument after `routes` is fixed for the endpoint's + # lifetime, so a launch passes these cached tensors and config values. self.partial = torch.zeros( config.ctas * config.logical_expert_count, dtype=torch.int32, @@ -335,51 +307,10 @@ def __init__(self, config: CudaSchedulerConfig, device: str | torch.device = "cu device=scheduler_device, ) self._connected = config.ep_size == 1 - self._plan_channel_lock = threading.Lock() - self._plan_channel_ptr = 0 - self._plan_channel_abi_version = 0 - self._plan_channel_words = 0 - self._plan_channel_route_features = 0 - self._plan_channel_bound = False - self._gpu_direct_bound = False - self._plan_channel_pending = False - self._plan_channel_stream_handle: int | None = None - self._plan_channel_torch_stream: torch.cuda.Stream | None = None - self._launch_stream_handle: int | None = None - - def _staging_stream(self) -> torch.cuda.Stream: - current = torch.cuda.current_stream(self.device) - bound = self._plan_channel_torch_stream - if bound is not None and _stream_handle(current) != self._plan_channel_stream_handle: - bound.wait_stream(current) - return bound - return current - - def stage(self, logical_expert_ids: torch.Tensor, *, validate_values: bool = False) -> None: - expected = (self.cfg.max_tokens_per_rank, self.cfg.topk) - if not isinstance(logical_expert_ids, torch.Tensor): - raise TypeError("logical_expert_ids must be a torch.Tensor") - if tuple(logical_expert_ids.shape) != expected: - raise ValueError(f"logical_expert_ids must have shape {expected}") - if logical_expert_ids.dtype != torch.int32 or not logical_expert_ids.is_contiguous(): - raise ValueError("logical_expert_ids must be contiguous torch.int32") - if not logical_expert_ids.is_cuda or logical_expert_ids.device != self.device: - raise ValueError("logical_expert_ids must be on the scheduler CUDA device") - if validate_values: - invalid = (logical_expert_ids < -1) | ( - logical_expert_ids >= self.cfg.logical_expert_count - ) - if bool(invalid.any().item()): - raise ValueError("logical expert IDs must be -1 or in [0,E)") - with torch.cuda.device(self.device): - stream = self._staging_stream() - with torch.cuda.stream(stream): - # No status clear: the kernel writes the failing generation into - # each flag rather than a bare 1, so a stale flag is - # distinguishable from a fresh one and clearing is only needed to - # acknowledge a fault (`reset_status()`). - self.logical_expert_ids.copy_(logical_expert_ids) - logical_expert_ids.record_stream(stream) + self._lock = threading.Lock() + # The bound copy's static stream: every fused launch shares the + # workspace and reusable barrier epochs, so all of them use it. + self._bound_stream: int | None = None def connect_peer_bases(self, peer_bases: object) -> None: if self._connected: @@ -396,134 +327,47 @@ def connect_peer_bases(self, peer_bases: object) -> None: torch.cuda.synchronize(self.device) self._connected = True - def bind_plan_channel( - self, - device_ptr: int, - stream: object, - *, - abi_version: int, - channel_words: int | None = None, - route_features: int = 0, - ) -> None: - with self._plan_channel_lock: - words = _plan_channel_word_count( - abi_version=abi_version, channel_words=channel_words, - route_features=route_features, ep=self.cfg.ep_size, - helpers=self.cfg.extra_slots_per_rank, - ) - if self._plan_channel_bound: - raise RuntimeError("plan channel binding is single-assignment") - if type(device_ptr) is not int or device_ptr <= 0 or device_ptr % 16: - raise ValueError("device_ptr must be a positive 16-byte aligned address") - handle = _stream_handle(stream) - with torch.cuda.device(self.device): - self._plan_channel_torch_stream = torch.cuda.ExternalStream( - handle, device=self.device - ) - self._plan_channel_ptr = device_ptr - self._plan_channel_abi_version = abi_version - self._plan_channel_words = words - self._plan_channel_route_features = route_features - self._plan_channel_stream_handle = handle - self._plan_channel_bound = True - def bind_gpu_direct(self, stream: object) -> None: - """Bind a copy consumer to the existing outputs/workspace, without publication. + """Bind the GPU-direct copy that launches this scheduler with its copy. - The consumer launches this scheduler with its copy on this static - stream through ``fused_submit_args``. No additional plan allocation or - GPU/host transfer is performed. + The copy reads the outputs and plan workspace in place, in the same + kernel, and every fused launch uses this static stream. No additional + plan allocation or GPU/host transfer is performed. """ - with self._plan_channel_lock: - if self._plan_channel_bound: - raise RuntimeError("plan channel binding is single-assignment") - handle = _stream_handle(stream) - if self._launch_stream_handle is not None and handle != self._launch_stream_handle: - raise RuntimeError("GPU-direct binding requires the existing static stream") - with torch.cuda.device(self.device): - self._plan_channel_torch_stream = torch.cuda.ExternalStream(handle, device=self.device) - self._plan_channel_ptr = 0 - self._plan_channel_abi_version = 0 - self._plan_channel_words = 0 - self._plan_channel_route_features = 0 - self._plan_channel_stream_handle = handle - self._plan_channel_bound = True - self._gpu_direct_bound = True - - def release_plan_channel(self) -> None: - """Allow another publication after its copy has been enqueued. - - GPU-direct readers and scheduler writers use the same static stream, - so stream order protects the single plan allocation without a host - completion wait. This does not release the weight bank's consumer lease. + with self._lock: + if self._bound_stream is not None: + raise RuntimeError("GPU-direct binding is single-assignment") + if not self._connected: + raise RuntimeError("EP peer bases are not connected") + self._bound_stream = _stream_handle(stream) + + def fused_submit_args( + self, routes: torch.Tensor, stream_handle: int, *, spin_cycles: int + ) -> tuple: + """Scheduler arguments of the bound copy's fused launch, ``routes`` + through ``valid_route_count`` of moe_rebalance_tma_submit_fused. + + ``routes`` is a caller-owned CUDA int32[T,K] tensor, read without input + staging. T may be zero and may differ across ranks, up to + max_tokens_per_rank. The kernel never reads outside those T rows and + writes -1 to the unused output tail. All ranks, including empty ranks, + must launch each generation in the same collective order. + + The caller must order input production before the bound static stream + and keep routes alive and unmodified until the scheduler completes. + Metadata checks below do not launch GPU work or synchronize; invalid + expert values are reported by the kernel through check_status(). + ``spin_cycles`` bounds this launch's cross-rank waits. """ - with self._plan_channel_lock: - if not self._plan_channel_bound or not self._plan_channel_pending: - raise RuntimeError("plan channel has no pending publication") - self._plan_channel_pending = False - - def launch( - self, stream: object | None = None, *, spin_cycles: int = SPIN_CYCLES - ) -> CudaSchedulerOutputs: _check_spin_cycles(spin_cycles) - if not self._connected: - raise RuntimeError("EP peer bases are not connected") - with self._plan_channel_lock: - if self._plan_channel_bound: - if self._plan_channel_pending: - raise RuntimeError("plan publication is still pending") - handle = self._plan_channel_stream_handle - if stream is not None and _stream_handle(stream) != handle: - raise RuntimeError("bound plan channel requires its static stream") - else: - handle = _stream_handle( - torch.cuda.current_stream(self.device) if stream is None else stream - ) - if self._launch_stream_handle is None: - self._launch_stream_handle = handle - elif handle != self._launch_stream_handle: - raise RuntimeError( - "CUDA scheduler uses one static stream because its workspace " - "and reusable barrier epochs are shared" - ) - self._launch_op( - self.logical_expert_ids, - int(handle), - valid_route_count=self.cfg.route_count, - spin_cycles=spin_cycles, + if self._bound_stream is None: + raise RuntimeError("a fused launch requires a bound GPU-direct copy") + if stream_handle != self._bound_stream: + raise RuntimeError( + "CUDA scheduler uses one static stream because its workspace " + "and reusable barrier epochs are shared" ) - if self._plan_channel_bound: - self._plan_channel_pending = True - return self.outputs - - def _launch_op( - self, - routes: torch.Tensor, - stream_handle: int, - *, - valid_route_count: int, - spin_cycles: int, - ) -> None: - stream = self._plan_channel_torch_stream - if stream is None or _stream_handle(stream) != stream_handle: - stream = torch.cuda.ExternalStream(stream_handle, device=self.device) - state = enter_cuda_stream(stream) - try: - torch.ops.trtllm.moe_rebalance_halo_q( - *self._halo_q_args( - routes, - valid_route_count, - enable_pdl=self.cfg.enable_pdl, - spin_cycles=spin_cycles, - ) - ) - finally: - exit_cuda_stream(state) - - def _halo_q_args( - self, routes: torch.Tensor, valid_route_count: int, *, enable_pdl: bool, spin_cycles: int - ) -> tuple: - """Positional moe_rebalance_halo_q arguments, shared by the fused copy op.""" + valid_tokens = self._check_routes(routes) return ( routes, self.outputs.physical_slot_ids, @@ -537,7 +381,6 @@ def _halo_q_args( self.grid_sync, self.plan_workspace, self.route_prefix, - self._plan_channel_ptr, self.cfg.ep_size, self.cfg.logical_expert_count, self.cfg.extra_slots_per_rank, @@ -545,99 +388,11 @@ def _halo_q_args( self.cfg.route_count, self.cfg.ctas, self.cfg.threads, - enable_pdl, spin_cycles, - self._plan_channel_abi_version, - self._plan_channel_words, - self._plan_channel_route_features, - valid_route_count, - ) - - def submit( - self, routes: torch.Tensor, stream_handle: int, *, - valid_tokens: int | None = None, - spin_cycles: int = SPIN_CYCLES, - ) -> CudaSchedulerOutputs: - """Read a caller-owned CUDA int32[T,K] tensor without input staging. - - T may be zero and may differ across ranks, up to max_tokens_per_rank. - By default all T rows are valid; valid_tokens can select a shorter - prefix of an existing allocation. The kernel never reads outside that - prefix and writes -1 to the unused output tail. All ranks, including - empty ranks, must submit each generation in the same collective order. - - The caller must order input production before this static stream and - keep routes alive and unmodified until the scheduler completes. This - includes allocator lifetime protection for temporary cross-stream - inputs. Metadata checks below do not launch GPU work or synchronize; - invalid expert values are reported by the kernel through check_status(). - - Graph capture fixes both the input pointer and valid_tokens. For replay - with changing pointers/lengths use eager submit; stage()+launch() remains - available for a fixed-capacity graph input buffer. - """ - _check_spin_cycles(spin_cycles) - with self._plan_channel_lock: - if self._plan_channel_bound: - if self._plan_channel_pending: - raise RuntimeError("plan publication is still pending") - if stream_handle != self._plan_channel_stream_handle: - raise RuntimeError("bound plan channel requires its static stream") - if self._launch_stream_handle is None: - self._launch_stream_handle = stream_handle - elif stream_handle != self._launch_stream_handle: - raise RuntimeError( - "CUDA scheduler uses one static stream because its workspace " - "and reusable barrier epochs are shared" - ) - self._submit_routes( - routes, stream_handle, valid_tokens=valid_tokens, spin_cycles=spin_cycles - ) - if self._plan_channel_bound: - self._plan_channel_pending = True - return self.outputs - - def fused_submit_args( - self, routes: torch.Tensor, stream_handle: int, *, spin_cycles: int - ) -> tuple: - """Scheduler arguments of the bound GPU-direct copy's fused launch. - - The input and stream contract is ``submit()``'s. The copy endpoint - launches this scheduler together with its copy, so no plan publication - becomes pending and the launch triggers no programmatic dependency. - ``spin_cycles`` bounds this launch's cross-rank waits. - """ - _check_spin_cycles(spin_cycles) - with self._plan_channel_lock: - if not self._gpu_direct_bound: - raise RuntimeError("a fused launch requires a bound GPU-direct copy") - if stream_handle != self._plan_channel_stream_handle: - raise RuntimeError("bound plan channel requires its static stream") - if self._launch_stream_handle is None: - self._launch_stream_handle = stream_handle - elif stream_handle != self._launch_stream_handle: - raise RuntimeError( - "CUDA scheduler uses one static stream because its workspace " - "and reusable barrier epochs are shared" - ) - valid_tokens = self._check_routes(routes, None) - return self._halo_q_args( - routes, valid_tokens * self.cfg.topk, enable_pdl=False, spin_cycles=spin_cycles - ) - - def _submit_routes( - self, routes: torch.Tensor, stream_handle: int, *, - valid_tokens: int | None, spin_cycles: int, - ) -> None: - valid_tokens = self._check_routes(routes, valid_tokens) - self._launch_op( - routes, - stream_handle, - valid_route_count=valid_tokens * self.cfg.topk, - spin_cycles=spin_cycles, + valid_tokens * self.cfg.topk, ) - def _check_routes(self, routes: torch.Tensor, valid_tokens: int | None) -> int: + def _check_routes(self, routes: torch.Tensor) -> int: if not isinstance(routes, torch.Tensor): raise TypeError("routes must be a torch.Tensor") if (routes.ndim != 2 or routes.shape[1] != self.cfg.topk @@ -647,30 +402,7 @@ def _check_routes(self, routes: torch.Tensor, valid_tokens: int | None) -> int: raise ValueError("routes must be contiguous torch.int32") if not routes.is_cuda or routes.device != self.device: raise ValueError("routes must be on the scheduler CUDA device") - if valid_tokens is None: - valid_tokens = routes.shape[0] - if type(valid_tokens) is not int or not 0 <= valid_tokens <= routes.shape[0]: - raise ValueError("valid_tokens must be an int in [0,T]") - return valid_tokens - - def run( - self, - logical_expert_ids: torch.Tensor, - stream: object | None = None, - *, - validate_values: bool = False, - spin_cycles: int = SPIN_CYCLES, - ) -> CudaSchedulerOutputs: - if stream is not None and not self._plan_channel_bound: - with torch.cuda.device(self.device): - external = torch.cuda.ExternalStream( - _stream_handle(stream), device=self.device - ) - with torch.cuda.stream(external): - self.stage(logical_expert_ids, validate_values=validate_values) - else: - self.stage(logical_expert_ids, validate_values=validate_values) - return self.launch(stream, spin_cycles=spin_cycles) + return routes.shape[0] def check_status(self) -> None: """Raise if any generation faulted since the last `reset_status()`. diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index 426644b96faa..eba948aeaf74 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -1460,41 +1460,23 @@ def test_scheduler_spin_bound_follows_warmup_transitions_and_autotuning( # AutoTuner profiling outside warmup gets the warmup bound too. tuning[0] = True assert bound.cycles() == warmup_cycles - # The startup warmup launch always uses the warmup bound. - assert bound.warmup_cycles == warmup_cycles -def test_scheduler_launches_pass_their_spin_bound_to_the_kernel(): - # Locate spin_cycles in the registered operator schema rather than hard-coding it. - op_source = ( - _ROOT / "cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.cpp" - ).read_text() - schema = "".join(re.findall(r'"([^"]*)"', op_source.split("m.def(", 1)[1].split(");", 1)[0])) - parameters = re.fullmatch(r"moe_rebalance_halo_q\((.*)\) -> \(\)", schema).group(1) +def test_fused_launch_passes_its_spin_bound_to_the_kernel(): + # Map the scheduler arguments onto the registered operator schema rather than hard-coding them. + op_source = (_ROOT / "cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp").read_text() + schemas = ["".join(re.findall(r'"([^"]*)"', block)) for block in op_source.split("m.def(")[1:]] + schema = next(text for text in schemas if text.startswith("moe_rebalance_tma_submit_fused(")) + parameters = re.fullmatch(r"moe_rebalance_tma_submit_fused\((.*)\) -> \(\)", schema).group(1) names = [parameter.split()[-1] for parameter in parameters.split(",")] - spin_index = names.index("spin_cycles") + scheduler_names = names[names.index("routes") :] - launched = [] - torch = SimpleNamespace( - Tensor=_Routes, - int32="int32", - ops=SimpleNamespace( - trtllm=SimpleNamespace(moe_rebalance_halo_q=lambda *args: launched.append(args)) - ), - ) namespace = _scheduler_spin_namespace() - namespace.update( - { - "torch": torch, - "enter_cuda_stream": lambda stream: stream, - "exit_cuda_stream": lambda state: None, - "_stream_handle": lambda stream: int(getattr(stream, "cuda_stream", stream)), - } - ) + namespace["torch"] = SimpleNamespace(Tensor=_Routes, int32="int32") scheduler_cls = _extract_class( _SCHEDULER_RUNTIME, "CudaPhysicalSlotScheduler", - {"launch", "_launch_op", "_halo_q_args", "fused_submit_args", "_check_routes"}, + {"fused_submit_args", "_check_routes"}, namespace, ) scheduler = scheduler_cls.__new__(scheduler_cls) @@ -1504,9 +1486,8 @@ def test_scheduler_launches_pass_their_spin_bound_to_the_kernel(): extra_slots_per_rank=4, local_rank=0, route_count=8192 * 6, - ctas=8, + ctas=4, threads=512, - enable_pdl=True, topk=6, max_tokens_per_rank=8192, ) @@ -1526,34 +1507,20 @@ def test_scheduler_launches_pass_their_spin_bound_to_the_kernel(): "route_prefix", ): setattr(scheduler, name, name) - scheduler.logical_expert_ids = "staged routes" scheduler.device = "cuda:0" - scheduler._connected = True - scheduler._plan_channel_lock = nullcontext() - scheduler._plan_channel_bound = False - scheduler._plan_channel_pending = False - scheduler._plan_channel_ptr = 0 - scheduler._plan_channel_abi_version = 0 - scheduler._plan_channel_words = 0 - scheduler._plan_channel_route_features = 0 - scheduler._plan_channel_torch_stream = SimpleNamespace(cuda_stream=19) - scheduler._launch_stream_handle = None - - # The standalone kernel keeps the native bound unless the caller passes one. - scheduler.launch(19) - scheduler.launch(19, spin_cycles=_WARMUP_SPIN_CYCLES) - assert [args[spin_index] for args in launched] == [_SERVING_SPIN_CYCLES, _WARMUP_SPIN_CYCLES] - - # The fused launch takes the same scheduler arguments with the caller's bound. - scheduler._gpu_direct_bound = True - scheduler._plan_channel_stream_handle = 19 + scheduler._bound_stream = 19 + routes = _Routes(16, []) routes.is_cuda = True args = scheduler.fused_submit_args(routes, 19, spin_cycles=_WARMUP_SPIN_CYCLES) - assert args[spin_index] == _WARMUP_SPIN_CYCLES - assert args[spin_index - 1] is False # No programmatic dependency in the fused kernel. + assert len(args) == len(scheduler_names) + assert args[scheduler_names.index("spin_cycles")] == _WARMUP_SPIN_CYCLES + assert args[scheduler_names.index("valid_route_count")] == 16 * 6 with pytest.raises(ValueError, match="positive exact int"): scheduler.fused_submit_args(routes, 19, spin_cycles=0) + # Every fused launch shares the workspace, so it must use the bound stream. + with pytest.raises(RuntimeError, match="one static stream"): + scheduler.fused_submit_args(routes, 20, spin_cycles=_WARMUP_SPIN_CYCLES) def test_raw_stream_switch_restores_the_previous_stream_and_device(): From 3dd8c3ec59ce02790922bb2906f69ebc9f76251e Mon Sep 17 00:00:00 2001 From: jintaop Date: Thu, 22 Oct 2026 20:12:25 -0700 Subject: [PATCH 33/59] [None][test] Follow main's executor API changes in the dynamic EPLB contract tests Main now passes release_nccl_window_owners to _release_cuda_graphs, routes PyExecutor.shutdown through _shutdown, and no longer takes an engine argument in worker_main or BaseWorker. Update the stubs and calls to match. Signed-off-by: jintaop --- .../_torch/moe/test_dynamic_eplb_contracts.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index eba948aeaf74..b6fbf281606c 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -633,7 +633,7 @@ def shutdown(self, *, local_safe): engine._cleanup_done = False engine.model_loader = None - engine._release_cuda_graphs = lambda: events.append("release_graphs") + engine._release_cuda_graphs = lambda **_: events.append("release_graphs") engine.moe_load_balancer = Manager() engine._runner = object() engine._model_caller = object() @@ -675,7 +675,7 @@ def shutdown(self, *, local_safe): engine = engine_cls() engine._cleanup_done = False engine.model_loader = None - engine._release_cuda_graphs = lambda: (_ for _ in ()).throw( + engine._release_cuda_graphs = lambda **_: (_ for _ in ()).throw( RuntimeError("graph cleanup failed") ) engine.moe_load_balancer = Manager() @@ -841,7 +841,7 @@ def async_worker_stop(self): raise RuntimeError("sampler stop failed") class Engine: - def _release_cuda_graphs(self): + def _release_cuda_graphs(self, **_): events.append("release_graphs") class DwdpManager: @@ -857,7 +857,7 @@ def __exit__(self, *args): executor_cls = _extract_class( _TORCH_ROOT / "pyexecutor/py_executor.py", "PyExecutor", - {"shutdown", "_finish_shutdown"}, + {"shutdown", "_shutdown", "_finish_shutdown"}, { "AsyncWorkerMixin": AsyncWorkerMixin, "torch": SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: False)), @@ -1149,7 +1149,7 @@ def __init__(self, **kwargs): ) args = worker_args() assert run( - lambda: worker_main(engine="engine", worker_queues=None, log_level="info", llm_args=args), + lambda: worker_main(worker_queues=None, log_level="info", llm_args=args), "mpi_comm", ) == expected + ["mpi_comm"] assert args.env_overrides == configured @@ -1172,9 +1172,7 @@ class GenerationExecutor: bases=("GenerationExecutor",), ) args = worker_args() - assert run(lambda: worker_cls(engine="engine", llm_args=args), "executor_init") == expected + [ - "executor_init" - ] + assert run(lambda: worker_cls(llm_args=args), "executor_init") == expected + ["executor_init"] assert args.env_overrides == configured # Ray worker: the first CUDA query follows launch-queue setup. From 6698251476eee9be73a31eacf2f02deb3895cfcb Mon Sep 17 00:00:00 2001 From: jintaop Date: Fri, 23 Oct 2026 01:27:48 -0700 Subject: [PATCH 34/59] [None][refactor] Move per-layer dynamic EPLB into its own package next to the load balancer Per-layer dynamic EPLB now lives in fused_moe/per_layer_dynamic_eplb/, next to moe_load_balancer.py, instead of one module beside it and two modules inside the MegaMoE backend package: - manager.py: PerLayerDynamicMoeLoadBalancer, the model-level owner. - layer.py: PerLayerDynamicMoeLayerLoadBalancer, the per-layer handle. - arena.py: the live-weight arena, moved from mega_moe/rebalance_live_arena.py. - planner.py: the HALO-Q and TMA submission, moved from mega_moe/rebalance_slot_scheduler.py. The import path of the two classes is unchanged. moe_load_balancer.py stays the only entry point, so fused_moe/__init__.py no longer re-exports them. Tests follow the new module paths. No behavior change. Signed-off-by: jintaop --- tensorrt_llm/_torch/moe/fused_moe/__init__.py | 4 - .../fused_moe/mega_moe/mega_moe_cute_dsl.py | 2 +- .../per_layer_dynamic_eplb/__init__.py | 17 + .../arena.py} | 0 .../layer.py} | 297 +----------------- .../per_layer_dynamic_eplb/manager.py | 285 +++++++++++++++++ .../planner.py} | 0 .../_torch/moe/test_dynamic_eplb_contracts.py | 14 +- tests/unittest/_torch/moe/test_moe_backend.py | 16 +- .../_torch/moe/test_per_layer_dynamic_eplb.py | 22 +- 10 files changed, 338 insertions(+), 319 deletions(-) create mode 100644 tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/__init__.py rename tensorrt_llm/_torch/moe/fused_moe/{mega_moe/rebalance_live_arena.py => per_layer_dynamic_eplb/arena.py} (100%) rename tensorrt_llm/_torch/moe/fused_moe/{per_layer_dynamic_eplb.py => per_layer_dynamic_eplb/layer.py} (52%) create mode 100644 tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/manager.py rename tensorrt_llm/_torch/moe/fused_moe/{mega_moe/rebalance_slot_scheduler.py => per_layer_dynamic_eplb/planner.py} (100%) diff --git a/tensorrt_llm/_torch/moe/fused_moe/__init__.py b/tensorrt_llm/_torch/moe/fused_moe/__init__.py index 2583c7c44fc1..2eade6470e55 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/__init__.py +++ b/tensorrt_llm/_torch/moe/fused_moe/__init__.py @@ -39,8 +39,6 @@ from .interface import MoE, MoEWeightLoadingMode from .moe_load_balancer import (MoeLoadBalancer, moe_load_balancer_set_repeated_for_next_layer) -from .per_layer_dynamic_eplb import (PerLayerDynamicMoeLayerLoadBalancer, - PerLayerDynamicMoeLoadBalancer) from .quantization import FusedMoEQuantScalesFP8 # yapf: disable from .routing import (BaseMoeRoutingMethod, DeepSeekV3MoeRoutingMethod, @@ -97,8 +95,6 @@ "MoEImplBase", "MoEImplClass", "MoeLoadBalancer", - "PerLayerDynamicMoeLayerLoadBalancer", - "PerLayerDynamicMoeLoadBalancer", "MoEWeightLoadingMode", "MiniMaxM2MoeRoutingMethod", "DeepSeekV4MoeRoutingMethod", diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py index 20796c3baa00..550cbe11306d 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py @@ -1142,7 +1142,7 @@ def create_weights(self): # Step 2-3: the same quant method registers every live Parameter. self.quant_method.create_weights(self) if self._rebalance_arena is not None: - from .rebalance_slot_scheduler import bind_live_weight_planes + from ..per_layer_dynamic_eplb.planner import bind_live_weight_planes bind_live_weight_planes(self.quant_method, self, self._rebalance_arena) # Step 4. diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/__init__.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/__init__.py new file mode 100644 index 000000000000..d37f18c5d3d6 --- /dev/null +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/__init__.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Model-owned lifecycle for per-layer dynamic expert load balancing. + +The standard EPLB implementation couples route updates to host weight +migration. Per-layer dynamic EPLB instead produces one device-side replica plan +for every MoE invocation and copies helper weights on a CUDA stream. It keeps +the existing load-balancer seams: ``manager`` owns the layers and their GPU +resources for the whole model, ``layer`` is the per-layer handle that plans each +route and releases its plan, ``arena`` allocates the live-weight storage, and +``planner`` submits HALO-Q and the TMA copy of its plan. +""" + +from .layer import PerLayerDynamicMoeLayerLoadBalancer +from .manager import PerLayerDynamicMoeLoadBalancer + +__all__ = ["PerLayerDynamicMoeLayerLoadBalancer", "PerLayerDynamicMoeLoadBalancer"] diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_live_arena.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/arena.py similarity index 100% rename from tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_live_arena.py rename to tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/arena.py diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/layer.py similarity index 52% rename from tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py rename to tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/layer.py index 7f3f53078c45..927098fc9688 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb.py +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/layer.py @@ -1,23 +1,18 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Model-owned lifecycle for per-layer dynamic expert load balancing. - -The standard EPLB implementation couples route updates to host weight -migration. Per-layer dynamic EPLB instead produces one device-side replica plan -for every MoE invocation and copies helper weights on a CUDA stream. These -classes preserve the existing load-balancer seams while keeping the scheduler, -arena, and their teardown under one model-level owner. -""" +"""Per-layer handle of per-layer dynamic EPLB: it plans each route and releases the plan.""" from __future__ import annotations import weakref -from typing import Any, Callable, Optional, Sequence +from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence import torch -from .impl_contract import MoEReplicaPlan -from .moe_load_balancer import BaseMoeLoadBalancer +from ..impl_contract import MoEReplicaPlan + +if TYPE_CHECKING: + from .manager import PerLayerDynamicMoeLoadBalancer class PerLayerDynamicMoeLayerLoadBalancer: @@ -179,7 +174,7 @@ def allocate_arena( rendezvous and before the kernel-plane Parameters are registered. ``bundle`` is the quant method's live-weight plane spec. """ - from .mega_moe.rebalance_live_arena import allocate_rebalance_arena + from .arena import allocate_rebalance_arena return allocate_rebalance_arena( home_experts=self.resident_experts_per_rank, @@ -218,7 +213,7 @@ def bind_backend(self, backend: Any) -> None: return backend_ref = weakref.ref(backend) - from .mega_moe.rebalance_slot_scheduler import build_rebalance_slot_scheduler_group + from .planner import build_rebalance_slot_scheduler_group if bool(getattr(backend, "tactic_autotune", False)): # Synthetic profiling may visit every helper row before a real @@ -356,279 +351,3 @@ def _clear_runtime(self) -> None: self._scheduler_group = None self._arena = None self._bound_backend_ref = None - - -class PerLayerDynamicMoeLoadBalancer(BaseMoeLoadBalancer): - """Model-level owner for per-layer dynamic EPLB layers and GPU resources.""" - - def __init__( - self, - *, - ep_rank: int, - ep_size: int, - config: Any, - mapping: Optional[Any] = None, - ) -> None: - # Per-layer dynamic EPLB has no C++ standard-EPLB object and creates no - # MPI split; its EP communicator is resolved by the first MoE layer. - if ep_size <= 1 or not 0 <= ep_rank < ep_size: - raise ValueError(f"invalid EPLB EP rank/size: rank={ep_rank}, size={ep_size}") - num_slots = int(config.num_slots) - if num_slots <= 0 or num_slots % ep_size != 0: - raise ValueError( - f"num_slots must be positive and divisible by EP size; got {num_slots}" - ) - - self.is_shutdown = False - self.ep_rank = int(ep_rank) - self.ep_size = int(ep_size) - self.config = config - self.mapping = mapping - self.num_slots = num_slots - self.num_local_slots = num_slots // ep_size - self.auxiliary_sms = int(getattr(config, "auxiliary_sms", 4)) - if self.auxiliary_sms <= 0: - raise ValueError("auxiliary_sms must be positive") - helper_slot_mode = getattr(config, "helper_slot_mode", "per_layer_resident") - if helper_slot_mode not in ("per_layer_resident", "shared_banks"): - raise ValueError(f"unknown helper_slot_mode {helper_slot_mode!r}") - # Resident helper experts need helper rows that only their layer writes. - self.helper_residency = helper_slot_mode == "per_layer_resident" - - # State read by the shared lifecycle and by the engine. - self._previous_balancer = None - self.single_layer_load_balancers: list[PerLayerDynamicMoeLayerLoadBalancer] = [] - self.next_layer_repeated_count: Optional[int] = None - self.iter_id = 0 - self.in_iter = False - self.enable_statistic = False - self.enable_update_weights = False - self._copy_stream: Optional[torch.cuda.Stream] = None - self._copy_stream_device: Optional[int] = None - self._ep_comm: Optional[Any] = None - self._scheduler_spin_bound: Optional[Any] = None - - # Providers use this manager as their scope owner. Resource references - # stay reachable here, never through Mapping or a process-global registry. - self._arena_providers: list[Any] = [] - self._rebalance_shared_slot_pools: dict[tuple, Any] = {} - self._rebalance_shared_slot_pool_order: list[Any] = [] - - def __del__(self) -> None: - # Teardown is collective and must be driven explicitly by the worker; - # a Python finalizer cannot safely enter EP collectives. - pass - - def __repr__(self) -> str: - return ( - f"PerLayerDynamicMoeLoadBalancer(ep_rank={self.ep_rank}, " - f"ep_size={self.ep_size}, num_slots={self.num_slots}, " - f"auxiliary_sms={self.auxiliary_sms})" - ) - - def is_static_routing(self) -> bool: - return False - - def set_repeated_for_next_layer(self, repeated_count: int) -> None: - if repeated_count <= 0: - raise ValueError("repeat count must be positive") - self.next_layer_repeated_count = int(repeated_count) - - def add_layer( - self, - expert_count: int, - top_k: int, - slot_count_per_rank: int, - aux_stream: Optional[torch.cuda.Stream] = None, - ) -> PerLayerDynamicMoeLayerLoadBalancer: - del aux_stream # The scheduler owns its high-priority copy stream. - self.config.validate_expert_capacity(expert_count) - if int(slot_count_per_rank) != self.num_local_slots: - raise ValueError( - "per-layer dynamic EPLB layer capacity must match the configured " - f"local slot count; got {slot_count_per_rank}, expected " - f"{self.num_local_slots}" - ) - repeated_count = self.next_layer_repeated_count or 1 - self.next_layer_repeated_count = None - layer = PerLayerDynamicMoeLayerLoadBalancer( - self, - layer_idx=len(self.single_layer_load_balancers), - expert_count=expert_count, - top_k=top_k, - compute_slots_per_rank=slot_count_per_rank, - repeated_count=repeated_count, - ) - self.single_layer_load_balancers.append(layer) - return layer - - def get_copy_stream(self, device: int) -> torch.cuda.Stream: - """Return this model's single highest-priority auxiliary stream.""" - device = int(device) - if self._copy_stream is None: - total_sms = int(torch.cuda.get_device_properties(device).multi_processor_count) - if self.auxiliary_sms >= total_sms: - raise ValueError( - "auxiliary_sms must leave at least one SM for the compute stream; " - f"got {self.auxiliary_sms} of {total_sms}" - ) - _, priority = torch.cuda.Stream.priority_range() - self._copy_stream = torch.cuda.Stream(device=device, priority=priority) - self._copy_stream_device = device - elif self._copy_stream_device != device: - raise RuntimeError("one per-layer dynamic EPLB manager cannot span CUDA devices") - return self._copy_stream - - def get_ep_comm(self, process_group: Any) -> Any: - """Return this model's adapter for its resolved EP ProcessGroup.""" - from .mega_moe.rebalance_live_arena import _EpComm - - if self._ep_comm is None: - self._ep_comm = _EpComm(process_group) - elif self._ep_comm.process_group is not process_group: - raise RuntimeError("one per-layer dynamic EPLB manager cannot span EP ProcessGroups") - self._ep_comm.validate_geometry(expected_rank=self.ep_rank, expected_world=self.ep_size) - return self._ep_comm - - def get_scheduler_spin_bound(self) -> Any: - """Return the HALO-Q spin bound shared by this model's layers. - - It is registered with the MoE communication timeout guard, which - lengthens it while the engine warms up. - """ - if self._scheduler_spin_bound is None: - from ...autotuner import AutoTuner - from .mega_moe.rebalance_slot_scheduler import SchedulerSpinBound - from .moe_comm_timeout_guard import ( - get_moe_comm_timeout_budgets, - register_moe_comm_timeout_proxy, - ) - - spin_bound = SchedulerSpinBound( - warmup_seconds=get_moe_comm_timeout_budgets().warmup_seconds, - is_tuning=lambda: AutoTuner.get().is_tuning_mode, - ) - register_moe_comm_timeout_proxy(spin_bound) - self._scheduler_spin_bound = spin_bound - return self._scheduler_spin_bound - - def create_arena_provider(self) -> Any: - """Create a per-layer provider scoped to this model manager.""" - from .mega_moe.rebalance_live_arena import SharedSlotArenaProvider - - provider = SharedSlotArenaProvider(self, per_layer_helper_sets=self.helper_residency) - self._arena_providers.append(provider) - return provider - - def adopt_arena_provider(self, provider: Any) -> None: - """Retain a transitional provider already created by the allocator.""" - if provider not in self._arena_providers: - self._arena_providers.append(provider) - - def _register_runtime(self, layer: PerLayerDynamicMoeLayerLoadBalancer) -> None: - if layer not in self.single_layer_load_balancers: - raise ValueError("cannot bind a layer owned by another EPLB manager") - assert layer._arena is not None - self.adopt_arena_provider(layer._arena.provider) - - def start_iter(self) -> None: - if self.in_iter: - raise RuntimeError("already in a MoE load-balancer iteration") - self.in_iter = True - - def end_iter(self) -> None: - if not self.in_iter: - raise RuntimeError("not in a MoE load-balancer iteration") - unfinished = [ - layer.layer_idx - for layer in self.single_layer_load_balancers - if layer.replica_plan is not None - ] - if unfinished: - raise RuntimeError(f"unfinished replica plans at iteration end: {unfinished}") - self.in_iter = False - self.iter_id += 1 - - def register_weight_slots_after_to_cuda(self) -> None: - # Live aliases are registered by the per-layer dynamic arena binder. - pass - - def finalize_model(self) -> None: - # Scheduler groups are built when each layer's final aliases exist. - pass - - def set_warm_up_iter_count(self, iter_count: int) -> None: - del iter_count - - @property - def has_submission_owner(self) -> bool: - return any(layer.has_submission_owner for layer in self.single_layer_load_balancers) - - def handoff_submission_owner(self, execution_stream_handle: int) -> None: - """Transfer all warmup-owned submitters after one caller device drain.""" - for layer in self.single_layer_load_balancers: - layer.handoff_submission_owner(execution_stream_handle) - - def release_quiesced_submission_owner(self, execution_stream_handle: int) -> None: - """Release executor ownership after the caller drains the device.""" - for layer in self.single_layer_load_balancers: - group = layer._scheduler_group - if group is not None: - group.release_quiesced_submission_owner(int(execution_stream_handle)) - - def shutdown(self, local_safe: bool = True) -> bool: - """Collectively close this model's pools in deterministic layer order.""" - if self.is_shutdown: - return True - - # Agree on local readiness before any rank starts destructive teardown. - # Every MoE layer resolves the EP communicator in create_weights, and a - # rank whose model load failed never owns an engine that reaches this - # method, so the communicator exists on every EP rank or on none. - local_safe = bool(local_safe and not self.in_iter) - comm = self._ep_comm - if comm is not None: - states = comm.allgather( - { - "local_safe": local_safe, - "pool_count": len(self._rebalance_shared_slot_pool_order), - } - ) - pool_counts = {state["pool_count"] for state in states} - if len(pool_counts) != 1: - raise RuntimeError( - "per-layer dynamic EPLB pool counts differ across the EP communicator" - ) - local_safe = all(state["local_safe"] for state in states) - if not local_safe: - return False - - # Snapshot creation order because each successful close removes itself. - for pool in list(self._rebalance_shared_slot_pool_order): - if not pool.close_collectively(local_safe=True): - return False - - # Pool teardown closes its registered producers. Calling close again is - # deliberately allowed and also covers a transitional group that has - # not yet been registered with a pool. - for layer in self.single_layer_load_balancers: - group = layer._scheduler_group - if group is not None: - group.close() - layer._clear_runtime() - self._arena_providers.clear() - self._ep_comm = None - if self._scheduler_spin_bound is not None: - from .moe_comm_timeout_guard import unregister_moe_comm_timeout_proxy - - unregister_moe_comm_timeout_proxy(self._scheduler_spin_bound) - self._scheduler_spin_bound = None - self._rebalance_shared_slot_pools.clear() - self._rebalance_shared_slot_pool_order.clear() - self._copy_stream = None - self._copy_stream_device = None - self.is_shutdown = True - return True - - -__all__ = ["PerLayerDynamicMoeLayerLoadBalancer", "PerLayerDynamicMoeLoadBalancer"] diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/manager.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/manager.py new file mode 100644 index 000000000000..156ed3b4cbda --- /dev/null +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/manager.py @@ -0,0 +1,285 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Model-level owner of the per-layer dynamic EPLB layers and their GPU resources.""" + +from __future__ import annotations + +from typing import Any, Optional + +import torch + +from ..moe_load_balancer import BaseMoeLoadBalancer +from .layer import PerLayerDynamicMoeLayerLoadBalancer + + +class PerLayerDynamicMoeLoadBalancer(BaseMoeLoadBalancer): + """Model-level owner for per-layer dynamic EPLB layers and GPU resources.""" + + def __init__( + self, + *, + ep_rank: int, + ep_size: int, + config: Any, + mapping: Optional[Any] = None, + ) -> None: + # Per-layer dynamic EPLB has no C++ standard-EPLB object and creates no + # MPI split; its EP communicator is resolved by the first MoE layer. + if ep_size <= 1 or not 0 <= ep_rank < ep_size: + raise ValueError(f"invalid EPLB EP rank/size: rank={ep_rank}, size={ep_size}") + num_slots = int(config.num_slots) + if num_slots <= 0 or num_slots % ep_size != 0: + raise ValueError( + f"num_slots must be positive and divisible by EP size; got {num_slots}" + ) + + self.is_shutdown = False + self.ep_rank = int(ep_rank) + self.ep_size = int(ep_size) + self.config = config + self.mapping = mapping + self.num_slots = num_slots + self.num_local_slots = num_slots // ep_size + self.auxiliary_sms = int(getattr(config, "auxiliary_sms", 4)) + if self.auxiliary_sms <= 0: + raise ValueError("auxiliary_sms must be positive") + helper_slot_mode = getattr(config, "helper_slot_mode", "per_layer_resident") + if helper_slot_mode not in ("per_layer_resident", "shared_banks"): + raise ValueError(f"unknown helper_slot_mode {helper_slot_mode!r}") + # Resident helper experts need helper rows that only their layer writes. + self.helper_residency = helper_slot_mode == "per_layer_resident" + + # State read by the shared lifecycle and by the engine. + self._previous_balancer = None + self.single_layer_load_balancers: list[PerLayerDynamicMoeLayerLoadBalancer] = [] + self.next_layer_repeated_count: Optional[int] = None + self.iter_id = 0 + self.in_iter = False + self.enable_statistic = False + self.enable_update_weights = False + self._copy_stream: Optional[torch.cuda.Stream] = None + self._copy_stream_device: Optional[int] = None + self._ep_comm: Optional[Any] = None + self._scheduler_spin_bound: Optional[Any] = None + + # Providers use this manager as their scope owner. Resource references + # stay reachable here, never through Mapping or a process-global registry. + self._arena_providers: list[Any] = [] + self._rebalance_shared_slot_pools: dict[tuple, Any] = {} + self._rebalance_shared_slot_pool_order: list[Any] = [] + + def __del__(self) -> None: + # Teardown is collective and must be driven explicitly by the worker; + # a Python finalizer cannot safely enter EP collectives. + pass + + def __repr__(self) -> str: + return ( + f"PerLayerDynamicMoeLoadBalancer(ep_rank={self.ep_rank}, " + f"ep_size={self.ep_size}, num_slots={self.num_slots}, " + f"auxiliary_sms={self.auxiliary_sms})" + ) + + def is_static_routing(self) -> bool: + return False + + def set_repeated_for_next_layer(self, repeated_count: int) -> None: + if repeated_count <= 0: + raise ValueError("repeat count must be positive") + self.next_layer_repeated_count = int(repeated_count) + + def add_layer( + self, + expert_count: int, + top_k: int, + slot_count_per_rank: int, + aux_stream: Optional[torch.cuda.Stream] = None, + ) -> PerLayerDynamicMoeLayerLoadBalancer: + del aux_stream # The scheduler owns its high-priority copy stream. + self.config.validate_expert_capacity(expert_count) + if int(slot_count_per_rank) != self.num_local_slots: + raise ValueError( + "per-layer dynamic EPLB layer capacity must match the configured " + f"local slot count; got {slot_count_per_rank}, expected " + f"{self.num_local_slots}" + ) + repeated_count = self.next_layer_repeated_count or 1 + self.next_layer_repeated_count = None + layer = PerLayerDynamicMoeLayerLoadBalancer( + self, + layer_idx=len(self.single_layer_load_balancers), + expert_count=expert_count, + top_k=top_k, + compute_slots_per_rank=slot_count_per_rank, + repeated_count=repeated_count, + ) + self.single_layer_load_balancers.append(layer) + return layer + + def get_copy_stream(self, device: int) -> torch.cuda.Stream: + """Return this model's single highest-priority auxiliary stream.""" + device = int(device) + if self._copy_stream is None: + total_sms = int(torch.cuda.get_device_properties(device).multi_processor_count) + if self.auxiliary_sms >= total_sms: + raise ValueError( + "auxiliary_sms must leave at least one SM for the compute stream; " + f"got {self.auxiliary_sms} of {total_sms}" + ) + _, priority = torch.cuda.Stream.priority_range() + self._copy_stream = torch.cuda.Stream(device=device, priority=priority) + self._copy_stream_device = device + elif self._copy_stream_device != device: + raise RuntimeError("one per-layer dynamic EPLB manager cannot span CUDA devices") + return self._copy_stream + + def get_ep_comm(self, process_group: Any) -> Any: + """Return this model's adapter for its resolved EP ProcessGroup.""" + from .arena import _EpComm + + if self._ep_comm is None: + self._ep_comm = _EpComm(process_group) + elif self._ep_comm.process_group is not process_group: + raise RuntimeError("one per-layer dynamic EPLB manager cannot span EP ProcessGroups") + self._ep_comm.validate_geometry(expected_rank=self.ep_rank, expected_world=self.ep_size) + return self._ep_comm + + def get_scheduler_spin_bound(self) -> Any: + """Return the HALO-Q spin bound shared by this model's layers. + + It is registered with the MoE communication timeout guard, which + lengthens it while the engine warms up. + """ + if self._scheduler_spin_bound is None: + from ....autotuner import AutoTuner + from ..moe_comm_timeout_guard import ( + get_moe_comm_timeout_budgets, + register_moe_comm_timeout_proxy, + ) + from .planner import SchedulerSpinBound + + spin_bound = SchedulerSpinBound( + warmup_seconds=get_moe_comm_timeout_budgets().warmup_seconds, + is_tuning=lambda: AutoTuner.get().is_tuning_mode, + ) + register_moe_comm_timeout_proxy(spin_bound) + self._scheduler_spin_bound = spin_bound + return self._scheduler_spin_bound + + def create_arena_provider(self) -> Any: + """Create a per-layer provider scoped to this model manager.""" + from .arena import SharedSlotArenaProvider + + provider = SharedSlotArenaProvider(self, per_layer_helper_sets=self.helper_residency) + self._arena_providers.append(provider) + return provider + + def adopt_arena_provider(self, provider: Any) -> None: + """Retain a transitional provider already created by the allocator.""" + if provider not in self._arena_providers: + self._arena_providers.append(provider) + + def _register_runtime(self, layer: PerLayerDynamicMoeLayerLoadBalancer) -> None: + if layer not in self.single_layer_load_balancers: + raise ValueError("cannot bind a layer owned by another EPLB manager") + assert layer._arena is not None + self.adopt_arena_provider(layer._arena.provider) + + def start_iter(self) -> None: + if self.in_iter: + raise RuntimeError("already in a MoE load-balancer iteration") + self.in_iter = True + + def end_iter(self) -> None: + if not self.in_iter: + raise RuntimeError("not in a MoE load-balancer iteration") + unfinished = [ + layer.layer_idx + for layer in self.single_layer_load_balancers + if layer.replica_plan is not None + ] + if unfinished: + raise RuntimeError(f"unfinished replica plans at iteration end: {unfinished}") + self.in_iter = False + self.iter_id += 1 + + def register_weight_slots_after_to_cuda(self) -> None: + # Live aliases are registered by the per-layer dynamic arena binder. + pass + + def finalize_model(self) -> None: + # Scheduler groups are built when each layer's final aliases exist. + pass + + def set_warm_up_iter_count(self, iter_count: int) -> None: + del iter_count + + @property + def has_submission_owner(self) -> bool: + return any(layer.has_submission_owner for layer in self.single_layer_load_balancers) + + def handoff_submission_owner(self, execution_stream_handle: int) -> None: + """Transfer all warmup-owned submitters after one caller device drain.""" + for layer in self.single_layer_load_balancers: + layer.handoff_submission_owner(execution_stream_handle) + + def release_quiesced_submission_owner(self, execution_stream_handle: int) -> None: + """Release executor ownership after the caller drains the device.""" + for layer in self.single_layer_load_balancers: + group = layer._scheduler_group + if group is not None: + group.release_quiesced_submission_owner(int(execution_stream_handle)) + + def shutdown(self, local_safe: bool = True) -> bool: + """Collectively close this model's pools in deterministic layer order.""" + if self.is_shutdown: + return True + + # Agree on local readiness before any rank starts destructive teardown. + # Every MoE layer resolves the EP communicator in create_weights, and a + # rank whose model load failed never owns an engine that reaches this + # method, so the communicator exists on every EP rank or on none. + local_safe = bool(local_safe and not self.in_iter) + comm = self._ep_comm + if comm is not None: + states = comm.allgather( + { + "local_safe": local_safe, + "pool_count": len(self._rebalance_shared_slot_pool_order), + } + ) + pool_counts = {state["pool_count"] for state in states} + if len(pool_counts) != 1: + raise RuntimeError( + "per-layer dynamic EPLB pool counts differ across the EP communicator" + ) + local_safe = all(state["local_safe"] for state in states) + if not local_safe: + return False + + # Snapshot creation order because each successful close removes itself. + for pool in list(self._rebalance_shared_slot_pool_order): + if not pool.close_collectively(local_safe=True): + return False + + # Pool teardown closes its registered producers. Calling close again is + # deliberately allowed and also covers a transitional group that has + # not yet been registered with a pool. + for layer in self.single_layer_load_balancers: + group = layer._scheduler_group + if group is not None: + group.close() + layer._clear_runtime() + self._arena_providers.clear() + self._ep_comm = None + if self._scheduler_spin_bound is not None: + from ..moe_comm_timeout_guard import unregister_moe_comm_timeout_proxy + + unregister_moe_comm_timeout_proxy(self._scheduler_spin_bound) + self._scheduler_spin_bound = None + self._rebalance_shared_slot_pools.clear() + self._rebalance_shared_slot_pool_order.clear() + self._copy_stream = None + self._copy_stream_device = None + self.is_shutdown = True + return True diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py similarity index 100% rename from tensorrt_llm/_torch/moe/fused_moe/mega_moe/rebalance_slot_scheduler.py rename to tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index b6fbf281606c..3ceaee1eadd7 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -88,7 +88,7 @@ def _extract_function(path: Path, name: str, namespace: dict): def test_shared_slot_provider_uses_model_owned_pool_state(): - source = _TORCH_ROOT / "moe/fused_moe/mega_moe/rebalance_live_arena.py" + source = _TORCH_ROOT / "moe/fused_moe/per_layer_dynamic_eplb/arena.py" provider_cls = _extract_class( source, "SharedSlotArenaProvider", @@ -116,7 +116,7 @@ def test_shared_slot_provider_uses_model_owned_pool_state(): def test_shared_slot_provider_pool_key_includes_complete_bundle(): - source = _TORCH_ROOT / "moe/fused_moe/mega_moe/rebalance_live_arena.py" + source = _TORCH_ROOT / "moe/fused_moe/per_layer_dynamic_eplb/arena.py" class FakePool: instances = [] @@ -182,7 +182,7 @@ class Bundle: @pytest.mark.parametrize("per_layer_helper_sets", [False, True]) def test_shared_slot_provider_pool_key_includes_helper_set_mode(per_layer_helper_sets): - source = _TORCH_ROOT / "moe/fused_moe/mega_moe/rebalance_live_arena.py" + source = _TORCH_ROOT / "moe/fused_moe/per_layer_dynamic_eplb/arena.py" class FakePool: instances = [] @@ -514,7 +514,7 @@ def test_vendored_live_plane_abi_matches_independent_goldens( def test_ep_collective_adapter_uses_the_resolved_process_group(): - source = _TORCH_ROOT / "moe/fused_moe/mega_moe/rebalance_live_arena.py" + source = _TORCH_ROOT / "moe/fused_moe/per_layer_dynamic_eplb/arena.py" tree = ast.parse(source.read_text()) selected = next( node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "_EpComm" @@ -1211,7 +1211,7 @@ def test_fused_submit_orders_generations_and_defers_the_route_wait(auxiliary_sms _cuda_getCurrentRawStream=lambda _: main.cuda_stream, ) torch = SimpleNamespace(cuda=cuda, _C=raw, Tensor=_Routes, int32="int32") - source = _TORCH_ROOT / "moe/fused_moe/mega_moe/rebalance_slot_scheduler.py" + source = _TORCH_ROOT / "moe/fused_moe/per_layer_dynamic_eplb/planner.py" namespace = {"torch": torch, "threading": SimpleNamespace(get_ident=lambda: state.thread)} group_cls = _extract_class(source, "RebalanceSlotSchedulerGroup", None, namespace) lease_cls = _extract_class(source, "_LiveBankLeaseProvider", None, namespace) @@ -1440,7 +1440,7 @@ def test_scheduler_spin_bound_follows_warmup_transitions_and_autotuning( runtime, ) bound_cls = _extract_class( - _TORCH_ROOT / "moe/fused_moe/mega_moe/rebalance_slot_scheduler.py", + _TORCH_ROOT / "moe/fused_moe/per_layer_dynamic_eplb/planner.py", "SchedulerSpinBound", None, {"Callable": Any, "Optional": Optional}, @@ -2068,7 +2068,7 @@ def test_on_autotune_builds_balanced_powerlaw_input(request): def test_per_layer_dynamic_shutdown_reports_in_iter_through_collective_preflight(): manager_cls = _extract_class( - _TORCH_ROOT / "moe/fused_moe/per_layer_dynamic_eplb.py", + _TORCH_ROOT / "moe/fused_moe/per_layer_dynamic_eplb/manager.py", "PerLayerDynamicMoeLoadBalancer", {"shutdown"}, {}, diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index 68cf1156ec07..4bd4ac797be0 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -124,10 +124,6 @@ MegaMoEDeepGemm, TrtllmCutedslMegaMoeNvfp4Impl, ) -from tensorrt_llm._torch.moe.fused_moe.mega_moe.rebalance_slot_scheduler import ( - RebalanceSlotSchedulerGroup, - _LiveBankLeaseProvider, -) from tensorrt_llm._torch.moe.fused_moe.moe_resolution import ( IMPL_PRIORITY, _reject_unsupported_activation, @@ -139,6 +135,10 @@ ExternalCommMoEScheduler, FusedCommMoEScheduler, ) +from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb.planner import ( + RebalanceSlotSchedulerGroup, + _LiveBankLeaseProvider, +) from tensorrt_llm._torch.moe.fused_moe.quantization import ( FusedMoEMethodBase, NVFP4FusedMoEMethod, @@ -2438,8 +2438,10 @@ def test_megamoe_cutedsl_mpi_bootstrap_keeps_executor_cuda_device(monkeypatch): def test_megamoe_cutedsl_cache_derived_state_reuses_rebalance_binding( monkeypatch, ): - from tensorrt_llm._torch.moe.fused_moe import PerLayerDynamicMoeLoadBalancer - from tensorrt_llm._torch.moe.fused_moe.mega_moe import rebalance_slot_scheduler + from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb import ( + PerLayerDynamicMoeLoadBalancer, + planner, + ) config = SimpleNamespace( num_slots=24, @@ -2473,7 +2475,7 @@ def test_megamoe_cutedsl_cache_derived_state_reuses_rebalance_binding( moe.layer_load_balancer = layer builder = MagicMock(return_value=SimpleNamespace(home_experts=4, helper_slots=2)) monkeypatch.setattr( - rebalance_slot_scheduler, + planner, "build_rebalance_slot_scheduler_group", builder, ) diff --git a/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py b/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py index 065f332911b6..eb0984380c04 100644 --- a/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py +++ b/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py @@ -8,9 +8,6 @@ import torch from tensorrt_llm._torch.moe.fused_moe.interface import MoE -from tensorrt_llm._torch.moe.fused_moe.mega_moe.rebalance_slot_scheduler import ( - RebalanceSlotSchedulerGroup, -) from tensorrt_llm._torch.moe.fused_moe.moe_load_balancer import ( BaseMoeLoadBalancer, MoeLoadBalancer, @@ -19,6 +16,9 @@ maybe_create_moe_load_balancer, ) from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb import PerLayerDynamicMoeLoadBalancer +from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb.planner import ( + RebalanceSlotSchedulerGroup, +) pytestmark = pytest.mark.cpu_only @@ -244,7 +244,7 @@ def test_layer_disables_standard_host_migration_and_checks_geometry(): def test_bind_backend_is_idempotent_before_scheduler_construction(monkeypatch): - from tensorrt_llm._torch.moe.fused_moe.mega_moe import rebalance_slot_scheduler + from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb import planner manager = _manager() layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) @@ -253,7 +253,7 @@ def test_bind_backend_is_idempotent_before_scheduler_construction(monkeypatch): group = _Group() builder = MagicMock(return_value=group) monkeypatch.setattr( - rebalance_slot_scheduler, + planner, "build_rebalance_slot_scheduler_group", builder, ) @@ -281,7 +281,7 @@ def test_bind_backend_is_idempotent_before_scheduler_construction(monkeypatch): def test_bind_backend_rejects_different_backend_or_arena_before_construction( monkeypatch, ): - from tensorrt_llm._torch.moe.fused_moe.mega_moe import rebalance_slot_scheduler + from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb import planner manager = _manager() layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) @@ -289,7 +289,7 @@ def test_bind_backend_rejects_different_backend_or_arena_before_construction( backend = _Backend(arena) builder = MagicMock(return_value=_Group()) monkeypatch.setattr( - rebalance_slot_scheduler, + planner, "build_rebalance_slot_scheduler_group", builder, ) @@ -367,7 +367,7 @@ def test_quiesced_owner_release_rejects_unfinished_generation(): def test_ep_communicator_is_model_owned(monkeypatch): - from tensorrt_llm._torch.moe.fused_moe.mega_moe import rebalance_live_arena + from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb import arena as arena_module creations = [] @@ -380,7 +380,7 @@ def __init__(self, process_group): def validate_geometry(self, *, expected_rank, expected_world): self.validations.append((expected_rank, expected_world)) - monkeypatch.setattr(rebalance_live_arena, "_EpComm", Comm) + monkeypatch.setattr(arena_module, "_EpComm", Comm) manager = _manager() process_group = object() @@ -395,7 +395,7 @@ def validate_geometry(self, *, expected_rank, expected_world): def test_layer_allocates_arena_from_model_owned_resources(monkeypatch): - from tensorrt_llm._torch.moe.fused_moe.mega_moe import rebalance_live_arena + from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb import arena as arena_module manager = _manager() layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) @@ -414,7 +414,7 @@ def create_arena_provider(): monkeypatch.setattr(manager, "get_ep_comm", get_ep_comm) monkeypatch.setattr(manager, "create_arena_provider", create_arena_provider) allocate = MagicMock(return_value=arena) - monkeypatch.setattr(rebalance_live_arena, "allocate_rebalance_arena", allocate) + monkeypatch.setattr(arena_module, "allocate_rebalance_arena", allocate) bundle, mapping = object(), object() allocated = layer.allocate_arena( From ca94215b44feeb245ac36637d38dd662ee0858ab Mon Sep 17 00:00:00 2001 From: jintaop Date: Fri, 23 Oct 2026 01:38:28 -0700 Subject: [PATCH 35/59] [None][refactor] Drive per-layer dynamic EPLB through the existing load-balancer seams Per-layer dynamic EPLB now runs its whole lifecycle through the seams standard EPLB already uses, instead of extra scheduler and executor hooks: - Route: FusedCommMoEScheduler calls _load_balancer_route once per chunk, on the same line as before. The planner consumes the logical routes that the layer-wise calibrator records or replays, so in that mode they are recorded before the route call; standard EPLB and EPLB-off still record the slots after it. - Release: the post-run_moe start_set_cpu_stage hook releases the chunk's replica plan on every chunk, while standard EPLB keeps its last-call gate. The scheduler's try/except around quantize and run_moe is gone. - Rollback: MoeLoadBalancerIterContext calls the new BaseMoeLoadBalancer.abort_iter when the forward raises. It defaults to end_iter, so standard EPLB is unchanged. The per-layer dynamic manager releases every pending plan, adds cleanup failures to the forward's error as notes (Python 3.10 has no add_note), and reports itself unsafe for collective shutdown while a plan is still pending. - Warmup: the planner binds the one MAIN execution stream rather than a CPU thread. Warmup and serving submit on that stream from different threads, so the warmup owner hand-off and the PyExecutor hooks that drove it are removed. The per-chunk CUDA work and its order are unchanged. Signed-off-by: jintaop --- .../_torch/moe/fused_moe/interface.py | 11 +- .../_torch/moe/fused_moe/moe_load_balancer.py | 14 +- .../_torch/moe/fused_moe/moe_scheduler.py | 118 ++--- .../fused_moe/per_layer_dynamic_eplb/layer.py | 39 +- .../per_layer_dynamic_eplb/manager.py | 55 +- .../per_layer_dynamic_eplb/planner.py | 67 +-- tensorrt_llm/_torch/pyexecutor/py_executor.py | 34 -- .../_torch/moe/test_dynamic_eplb_contracts.py | 476 ++++++++++++++---- .../_torch/moe/test_per_layer_dynamic_eplb.py | 86 ++-- 9 files changed, 548 insertions(+), 352 deletions(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/interface.py b/tensorrt_llm/_torch/moe/fused_moe/interface.py index 07dea4102aed..b13c06dbd262 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/interface.py +++ b/tensorrt_llm/_torch/moe/fused_moe/interface.py @@ -779,15 +779,12 @@ def _load_balancer_replica_plan(self): return None return self.layer_load_balancer.require_replica_plan() - def _load_balancer_finish_replica_plan(self, plan) -> None: - """Release a caller-owned plan after its last MAIN-stream consumer.""" - if plan is None or not self.layer_load_balancer: - return - self.layer_load_balancer.finish_replica_plan(plan) - def _load_balancer_start_set_cpu_stage(self, is_last_call: bool): """Start CPU stage in load balancer.""" - if self._using_dynamic_load_balancer() and is_last_call: + # A per-layer dynamic EPLB layer plans every chunk and releases that + # chunk's replica plan here, after the chunk's MoE kernel. + if self._load_balancer_uses_replica_plan() or ( + self._using_dynamic_load_balancer() and is_last_call): self.layer_load_balancer.start_set_cpu_stage() def _load_balancer_done_set_cpu_stage(self, is_last_call: bool): diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py b/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py index 83467711a693..a8f48b9b3062 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py @@ -775,8 +775,9 @@ class BaseMoeLoadBalancer(ABC): Model construction enters the balancer as a context manager so that MoE layers register through ``get_moe_load_balancer``. The engine drives each forward through ``MoeLoadBalancerIterContext`` (``set_iter_info``, - ``start_iter`` and ``end_iter``), and workers release resources with a - collective ``shutdown``. Subclasses initialize ``_previous_balancer``, + ``start_iter``, and ``end_iter`` or, when the forward raises, + ``abort_iter``), and workers release resources with a collective + ``shutdown``. Subclasses initialize ``_previous_balancer``, ``enable_statistic`` and ``enable_update_weights``. """ @@ -819,6 +820,10 @@ def start_iter(self) -> None: def end_iter(self) -> None: """End the current forward iteration.""" + def abort_iter(self, error: BaseException) -> None: + """End an iteration whose forward raised ``error``; defaults to ``end_iter``.""" + self.end_iter() + @abstractmethod def shutdown(self, local_safe: bool = True) -> bool: """Collectively release resources; return False if any rank is unsafe.""" @@ -1205,7 +1210,10 @@ def __exit__(self, exc_type, exc_val, exc_tb): """ if self.moe_load_balancer is not None and not self.moe_load_balancer.is_static_routing( ): - self.moe_load_balancer.end_iter() + if exc_type is None: + self.moe_load_balancer.end_iter() + else: + self.moe_load_balancer.abort_iter(exc_val) return False diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py index a01537dda4f6..f755e214b8dc 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py @@ -83,6 +83,19 @@ def _split_rows(t: torch.Tensor, sizes: List[int]) -> List[torch.Tensor]: return list(t.split(sizes)) +def _record_routes( + moe: "ConfigurableMoE", routes: Optional[torch.Tensor] +) -> Optional[torch.Tensor]: + """Record ``routes`` for expert statistics and the layer-wise calibrator. + + Returns ``routes``, or the routes the calibrator replays in their place. + """ + if routes is not None: + ExpertStatistic.set_layer(moe.layer_idx) + ExpertStatistic.maybe_add_info(moe.num_slots, routes) + return get_calibrator().maybe_collect_or_replay_slots(moe.num_slots, routes) + + def _wait_input_ready(event: Optional[torch.cuda.Event]) -> None: """Join a pending ``input_ready_event`` before the first read of ``x``. @@ -1551,8 +1564,10 @@ def _forward_chunk( ) # ----- EPLB: update stats + remap expert ids -> slot ids ----- + # Per-layer dynamic EPLB plans from the routes that the layer-wise + # calibrator records or replays, so they are recorded before the route; + # otherwise the recorded routes are the slots routed to. uses_replica_plan = moe._load_balancer_uses_replica_plan() - replica_plan = None if moe.layer_load_balancer: moe._load_balancer_done_wait_gpu_stage(is_first_call) # ignore_allreduce=False: the fused kernel has no side channel @@ -1564,68 +1579,53 @@ def _forward_chunk( is_last_call, ignore_allreduce=False, ) - # Standard EPLB remaps here. Per-layer dynamic EPLB must consume the - # final logical routes after calibration/replay below. - token_selected_slots = ( - token_selected_experts - if uses_replica_plan - else moe._load_balancer_route(token_selected_experts, moe.use_dp) - ) + if uses_replica_plan: + token_selected_experts = _record_routes(moe, token_selected_experts) + token_selected_slots = moe._load_balancer_route(token_selected_experts, moe.use_dp) else: token_selected_slots = token_selected_experts - if token_selected_slots is not None: - ExpertStatistic.set_layer(moe.layer_idx) - ExpertStatistic.maybe_add_info(moe.num_slots, token_selected_slots) - token_selected_slots = get_calibrator().maybe_collect_or_replay_slots( - moe.num_slots, token_selected_slots - ) - - if uses_replica_plan: - token_selected_slots = moe._load_balancer_route(token_selected_slots, moe.use_dp) - replica_plan = moe._load_balancer_replica_plan() - - try: - # ----- quantize / prepare ----- - if getattr(moe.backend, "supports_fused_prepare", lambda: False)(): - # MegaMoE can fuse BF16->MXFP8 quantization with SymmBuffer - # topk copies, so keep activations for run_moe preparation. - moe_input = x_chunk_real - x_sf = None - else: - # Each fused-comm backend owns its empty-tensor layout. - moe_input, x_sf = moe.backend.quantize_input(x_chunk_real) - - # Select one rank-identical adaptive bucket on every EP rank. - set_adaptive = getattr(moe.backend, "set_adaptive_launch_tokens", None) - if set_adaptive is not None: - set_adaptive(max(all_rank_num_tokens) if all_rank_num_tokens else None) - - # ----- MoE compute ----- - # The manager owns the plan; the backend only consumes its - # geometry, READY terminals and late route dependency. - out = moe.backend.run_moe( - MoERunContext( - token_selected_experts=token_selected_slots, - token_final_scales=token_final_scales, - x=moe_input, - x_sf=x_sf, - output_dtype=output_dtype, - replica_plan=replica_plan, - ) - ) - except BaseException as error: # noqa: BLE001 - if replica_plan is not None: - try: - moe._load_balancer_finish_replica_plan(replica_plan) - except BaseException as cleanup_error: # noqa: BLE001 - error.add_note( - "Failed to release the per-layer dynamic EPLB plan after " - f"forward failure: {cleanup_error!r}" - ) - raise + if not uses_replica_plan: + token_selected_slots = _record_routes(moe, token_selected_slots) + replica_plan = moe._load_balancer_replica_plan() if uses_replica_plan else None + + # ----- quantize / prepare ----- + if getattr(moe.backend, "supports_fused_prepare", lambda: False)(): + # MegaMoE can fuse BF16->MXFP8 quantization with the SymmBuffer + # topk copies, so keep the original activations and let run_moe + # prepare its workspace. + moe_input = x_chunk_real + x_sf = None else: - moe._load_balancer_finish_replica_plan(replica_plan) + # Delegate to ``backend.quantize_input`` so each fused-comm backend + # owns its own empty-tensor layout. Both MegaMoEDeepGemm and + # MegaMoECuteDsl short-circuit ``x.shape[0] == 0`` inside their + # quantize_input contracts. + moe_input, x_sf = moe.backend.quantize_input(x_chunk_real) + + # CuteDSL needs the scheduler's rank-identical chunk maximum to select + # one adaptive bucket on every EP rank; using a local token count could + # diverge and deadlock its in-kernel NVLink barrier. + set_adaptive = getattr(moe.backend, "set_adaptive_launch_tokens", None) + if set_adaptive is not None: + set_adaptive(max(all_rank_num_tokens) if all_rank_num_tokens else None) + + # ----- MoE compute ----- + # ``token_selected_slots`` is in [0, num_slots), matching the kernel's + # ``num_experts`` template parameter (SymmBuffer / weights sized to + # num_slots in quantization.py). + # Fused-comm backends own the EP exchange, so there is no comm plan: + # nothing outside the fused kernel decided anything about this forward. + out = moe.backend.run_moe( + MoERunContext( + token_selected_experts=token_selected_slots, + token_final_scales=token_final_scales, + x=moe_input, + x_sf=x_sf, + output_dtype=output_dtype, + replica_plan=replica_plan, + ) + ) # ----- EPLB: start/done CPU rebalance, AFTER run_moe ----- # The external-comm path overlaps CPU stage with ``comm.combine``; diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/layer.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/layer.py index 927098fc9688..4b5607baff38 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/layer.py +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/layer.py @@ -15,6 +15,15 @@ from .manager import PerLayerDynamicMoeLoadBalancer +def _append_note(error: BaseException, note: str) -> None: + """Attach ``note`` to ``error`` like ``BaseException.add_note``, which Python 3.10 lacks.""" + notes = getattr(error, "__notes__", None) + if type(notes) is not list: + notes = [] + error.__notes__ = notes + notes.append(note) + + class PerLayerDynamicMoeLayerLoadBalancer: """Adapter between the common MoE routing seam and one replica planner.""" @@ -100,18 +109,19 @@ def set_initial_weight_assignments( "valid in standard mode." ) - # These hooks are part of the standard EPLB layer interface. Routing and - # statistics are produced by HALO-Q, so the corresponding CPU stages are - # intentionally empty here. + def start_set_cpu_stage(self) -> None: + """Release the plan of the chunk whose MoE kernel was just enqueued.""" + self.finish_replica_plan() + + # The other hooks of the standard EPLB layer interface are intentionally + # empty: HALO-Q produces the routing and statistics on the GPU, and no CPU + # stage migrates weights. def start_wait_gpu_stage(self) -> None: pass def done_wait_gpu_stage(self) -> None: pass - def start_set_cpu_stage(self) -> None: - pass - def done_set_cpu_stage(self) -> None: pass @@ -287,11 +297,7 @@ def route( "rebalance generation cleanup failed and the group " f"requires teardown: {cleanup_error!r}" ) - notes = getattr(error, "__notes__", None) - if type(notes) is not list: - notes = [] - error.__notes__ = notes - notes.append("overlap_fn failed after HALO-Q/TMA enqueue; " + cleanup_note) + _append_note(error, "overlap_fn failed after HALO-Q/TMA enqueue; " + cleanup_note) raise physical, generation = self._scheduler_group.plan_finish(part, defer_wait=True) generation = int(generation) @@ -335,17 +341,6 @@ def finish_replica_plan( self._scheduler_group.finish(completion_event=completion_event) self._pending_replica_plan = None - @property - def has_submission_owner(self) -> bool: - group = self._scheduler_group - return bool(group is not None and group.has_submission_owner) - - def handoff_submission_owner(self, execution_stream_handle: int) -> None: - """Release this layer's warmup submitter after a caller device drain.""" - group = self._scheduler_group - if group is not None and group.has_submission_owner: - group.release_warmup_owner_after_sync(int(execution_stream_handle)) - def _clear_runtime(self) -> None: self._pending_replica_plan = None self._scheduler_group = None diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/manager.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/manager.py index 156ed3b4cbda..61d1fb288c09 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/manager.py +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/manager.py @@ -9,7 +9,7 @@ import torch from ..moe_load_balancer import BaseMoeLoadBalancer -from .layer import PerLayerDynamicMoeLayerLoadBalancer +from .layer import PerLayerDynamicMoeLayerLoadBalancer, _append_note class PerLayerDynamicMoeLoadBalancer(BaseMoeLoadBalancer): @@ -203,6 +203,28 @@ def end_iter(self) -> None: self.in_iter = False self.iter_id += 1 + def abort_iter(self, error: BaseException) -> None: + """End a failed iteration after releasing the plans it left unfinished. + + A forward that raises between a layer's route and its post-``run_moe`` + hook leaves that layer's plan pending. + """ + try: + for layer in self.single_layer_load_balancers: + if layer.replica_plan is None: + continue + try: + layer.finish_replica_plan() + except BaseException as cleanup_error: # noqa: BLE001 + _append_note( + error, + "per-layer dynamic EPLB could not release the plan of layer " + f"{layer.layer_idx}: {cleanup_error!r}", + ) + finally: + self.in_iter = False + self.iter_id += 1 + def register_weight_slots_after_to_cuda(self) -> None: # Live aliases are registered by the per-layer dynamic arena binder. pass @@ -214,32 +236,21 @@ def finalize_model(self) -> None: def set_warm_up_iter_count(self, iter_count: int) -> None: del iter_count - @property - def has_submission_owner(self) -> bool: - return any(layer.has_submission_owner for layer in self.single_layer_load_balancers) - - def handoff_submission_owner(self, execution_stream_handle: int) -> None: - """Transfer all warmup-owned submitters after one caller device drain.""" - for layer in self.single_layer_load_balancers: - layer.handoff_submission_owner(execution_stream_handle) - - def release_quiesced_submission_owner(self, execution_stream_handle: int) -> None: - """Release executor ownership after the caller drains the device.""" - for layer in self.single_layer_load_balancers: - group = layer._scheduler_group - if group is not None: - group.release_quiesced_submission_owner(int(execution_stream_handle)) - def shutdown(self, local_safe: bool = True) -> bool: """Collectively close this model's pools in deterministic layer order.""" if self.is_shutdown: return True - # Agree on local readiness before any rank starts destructive teardown. - # Every MoE layer resolves the EP communicator in create_weights, and a - # rank whose model load failed never owns an engine that reaches this - # method, so the communicator exists on every EP rank or on none. - local_safe = bool(local_safe and not self.in_iter) + # Agree on local readiness before any rank starts destructive teardown; + # a rank is unsafe while any plan is pending. Every MoE layer resolves + # the EP communicator in create_weights, and a rank whose model load + # failed never owns an engine that reaches this method, so the + # communicator exists on every EP rank or on none. + local_safe = bool( + local_safe + and not self.in_iter + and all(layer.replica_plan is None for layer in self.single_layer_load_balancers) + ) comm = self._ep_comm if comm is not None: states = comm.allgather( diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py index 2a25de9e95d5..b2fac04d5051 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Single-main-thread HALO-Q and in-switch TMA weight-copy integration. +"""HALO-Q and in-switch TMA weight-copy integration on one MAIN stream. One highest-priority stream is shared across layers. plan_schedule() enqueues one kernel that runs HALO-Q and then the TMA copy of its plan on the auxiliary @@ -20,7 +20,6 @@ from __future__ import annotations -import threading from typing import Any, Callable, Optional, Tuple import torch @@ -167,7 +166,6 @@ def __init__( self._route_wait_generation = 0 # Routes of the latest launch, held until MAIN waits for its route epoch. self._routes_in_flight: Optional[torch.Tensor] = None - self._owner_thread_id: Optional[int] = None self._execution_stream_handle: Optional[int] = None self._closed = False self.scheduler = None @@ -288,7 +286,6 @@ def close(self) -> None: if self.symmetric is not None: self.symmetric.close() self._closed = True - self._owner_thread_id = None self._execution_stream_handle = None self._plan_part = None self._pending_release = None @@ -299,74 +296,26 @@ def close(self) -> None: self.symmetric = None self.arena = None - @property - def has_submission_owner(self) -> bool: - """Whether a CPU thread currently owns MAIN submissions.""" - return self._owner_thread_id is not None - def _check_owner(self) -> None: if self._closed: raise RuntimeError("Rebalance scheduler group is closed") - thread_id = threading.get_ident() # Runs several times per layer: read the raw device and stream instead - # of building a torch.cuda.Stream, and return early for the bound owner. + # of building a torch.cuda.Stream, and return early for the bound MAIN. if torch._C._cuda_getDevice() != self.device: raise RuntimeError("Rebalance submission requires its bound CUDA device") handle = torch._C._cuda_getCurrentRawStream(self.device) - if thread_id == self._owner_thread_id and handle == self._execution_stream_handle: + if handle == self._execution_stream_handle: return - if self._owner_thread_id is not None: - raise RuntimeError("Rebalance must be submitted by one MAIN thread and stream") - if self._execution_stream_handle is not None and handle != self._execution_stream_handle: - raise RuntimeError("Rebalance MAIN stream cannot change during owner handoff") + if self._execution_stream_handle is not None: + raise RuntimeError( + "Rebalance MAIN stream cannot change; every executor of a model " + "engine must submit on its one execution stream" + ) stream = torch.cuda.current_stream(self.device) if handle == self._stream_handle or stream.priority <= self.copy_stream.priority: raise RuntimeError("MAIN must be distinct from the higher-priority COPY stream") - self._owner_thread_id = thread_id self._execution_stream_handle = handle - def release_warmup_owner_after_sync(self, execution_stream_handle: int) -> None: - """Release the startup CPU owner after the caller drains this device. - - Keep the fixed MAIN stream, generation counters, and pending lease - release. The first serving submission binds its CPU owner normally. - """ - if self._closed: - raise RuntimeError("Rebalance scheduler group is closed") - if self._owner_thread_id != threading.get_ident(): - raise RuntimeError("Only the warmup submitter may release its owner") - if self._execution_stream_handle != int(execution_stream_handle): - raise RuntimeError("Warmup owner handoff requires the bound MAIN stream") - if ( - self._plan_part is not None - or self._generation != self._finished_generation - or self._generation != self.plan_calls - ): - raise RuntimeError("Warmup owner handoff requires all plans to be finished") - self._owner_thread_id = None - - def release_quiesced_submission_owner(self, execution_stream_handle: int) -> None: - """Release a stopped executor's MAIN owner after a device drain. - - Unlike the warmup handoff, this is called by the shutdown thread after - the serving worker has exited. The caller must synchronize the device - first, so neither the old CPU thread nor its MAIN stream can still be - submitting work. - """ - if self._closed: - raise RuntimeError("Rebalance scheduler group is closed") - handle = int(execution_stream_handle) - if ( - self._plan_part is not None - or self._generation != self._finished_generation - or self._generation != self.plan_calls - ): - raise RuntimeError("Quiesced owner release requires all plans to be finished") - if self._execution_stream_handle is not None and self._execution_stream_handle != handle: - raise RuntimeError("Quiesced owner release requires the bound MAIN stream") - self._owner_thread_id = None - self._execution_stream_handle = None - def _record_event(self, event_handle: int, stream_handle: int) -> None: (error,) = self._driver.cuEventRecord(event_handle, stream_handle) if error != self._driver.CUresult.CUDA_SUCCESS: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 976b027d8f90..356652978015 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1122,8 +1122,6 @@ def _initialize_execution(self, resources: PyExecutorResources) -> None: self.encoder_launch_executor.submit( self._warmup_encoder_cuda_graphs_enc_dec).result() - self._handoff_rebalance_warmup_owners() - # Snapshot some cumulative KV cache counters so that stats reported to # users exclude blocks reused and missed during warmup dummy requests. if hasattr(self.kv_cache_manager, 'snapshot_warmup_baseline'): @@ -1402,38 +1400,6 @@ def _kv_connector_moves_kv(self) -> bool: connector = getattr(self, "kv_connector_manager", None) return connector is not None and not connector.capacity_only - def _handoff_rebalance_warmup_owners(self) -> None: - """Transfer model-level EPLB submitters from warmup to serving.""" - managers = {} - for engine in (self.model_engine, self.draft_model_engine): - manager = getattr(engine, "moe_load_balancer", None) - handoff = getattr(manager, "handoff_submission_owner", None) - if callable(handoff): - managers[id(manager)] = manager - if not managers: - return - - # Target, draft, and encoder warmup have all completed. One device - # drain covers every manager's compute and copy streams before the - # serving worker binds submission ownership. - torch.cuda.synchronize(self.execution_stream.device) - execution_stream_handle = int(self.execution_stream.cuda_stream) - for manager in managers.values(): - manager.handoff_submission_owner(execution_stream_handle) - - def _release_quiesced_rebalance_submission_owners(self) -> None: - """Release EPLB submitters after the worker exits and device drains.""" - execution_stream_handle = int(self.execution_stream.cuda_stream) - managers = {} - for engine in (self.model_engine, self.draft_model_engine): - manager = getattr(engine, "moe_load_balancer", None) - release = getattr(manager, "release_quiesced_submission_owner", - None) - if callable(release): - managers[id(manager)] = manager - for manager in managers.values(): - manager.release_quiesced_submission_owner(execution_stream_handle) - def _maybe_init_kv_connector_manager(self) -> None: if self.kv_connector_manager is not None: # A connector that moves no KV cannot contend with the transceiver diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index 3ceaee1eadd7..8f3e123389ed 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -12,6 +12,7 @@ import os import re import sys +import threading from contextlib import nullcontext from pathlib import Path from types import ModuleType, SimpleNamespace @@ -1198,7 +1199,7 @@ class GenerationExecutor: @pytest.mark.parametrize("auxiliary_sms", [4, 8]) def test_fused_submit_orders_generations_and_defers_the_route_wait(auxiliary_sms): trace = [] - state = SimpleNamespace(device=0, thread=7) + state = SimpleNamespace(device=0) main = SimpleNamespace(cuda_stream=44, priority=0) copy = SimpleNamespace(cuda_stream=19, priority=-1) cuda = SimpleNamespace( @@ -1212,7 +1213,7 @@ def test_fused_submit_orders_generations_and_defers_the_route_wait(auxiliary_sms ) torch = SimpleNamespace(cuda=cuda, _C=raw, Tensor=_Routes, int32="int32") source = _TORCH_ROOT / "moe/fused_moe/per_layer_dynamic_eplb/planner.py" - namespace = {"torch": torch, "threading": SimpleNamespace(get_ident=lambda: state.thread)} + namespace = {"torch": torch} group_cls = _extract_class(source, "RebalanceSlotSchedulerGroup", None, namespace) lease_cls = _extract_class(source, "_LiveBankLeaseProvider", None, namespace) group = group_cls.__new__(group_cls) @@ -1229,7 +1230,6 @@ def test_fused_submit_orders_generations_and_defers_the_route_wait(auxiliary_sms group._pending_release = None group._routes_in_flight = None group._closed = False - group._owner_thread_id = None group._execution_stream_handle = None group.device = 0 group.topk = 6 @@ -1310,13 +1310,9 @@ def submit_fused(result, routes, epoch, target, *, spin_cycles): ) assert launch_bounds == [_WARMUP_SPIN_CYCLES, _SERVING_SPIN_CYCLES] - # The bound-owner fast path still rejects another thread, stream or device. - state.thread = 8 - with pytest.raises(RuntimeError, match="one MAIN thread and stream"): - group.plan_schedule(_Routes(8, trace)) - state.thread = 7 + # The bound-stream fast path still rejects another stream or device. main.cuda_stream = 45 - with pytest.raises(RuntimeError, match="one MAIN thread and stream"): + with pytest.raises(RuntimeError, match="MAIN stream cannot change"): group.plan_schedule(_Routes(8, trace)) main.cuda_stream = 44 state.device = 1 @@ -1356,6 +1352,64 @@ def submit_fused(result, routes, epoch, target, *, spin_cycles): group.wait_for_routes() +def test_rebalance_submission_binds_one_main_stream_for_every_thread(): + priorities = {19: -1, 44: 0, 45: 0, 46: -1} + current = SimpleNamespace(stream=46) + torch = SimpleNamespace( + _C=SimpleNamespace( + _cuda_getDevice=lambda: 0, + _cuda_getCurrentRawStream=lambda _: current.stream, + ), + cuda=SimpleNamespace( + current_stream=lambda _: SimpleNamespace( + cuda_stream=current.stream, priority=priorities[current.stream] + ) + ), + ) + group_cls = _extract_class( + _TORCH_ROOT / "moe/fused_moe/per_layer_dynamic_eplb/planner.py", + "RebalanceSlotSchedulerGroup", + {"_check_owner"}, + {"torch": torch}, + ) + group = group_cls() + group._closed = False + group.device = 0 + group.copy_stream = SimpleNamespace(cuda_stream=19, priority=-1) + group._stream_handle = 19 + group._execution_stream_handle = None + + # Neither COPY nor a stream of its priority can become MAIN. + for stream in (46, 19): + current.stream = stream + with pytest.raises(RuntimeError, match="distinct from the higher-priority COPY"): + group._check_owner() + assert group._execution_stream_handle is None + + # The first other stream binds as MAIN, and any thread may submit on it: + # warmup and serving run on different threads but share one stream. + current.stream = 44 + group._check_owner() + assert group._execution_stream_handle == 44 + errors = [] + + def submit(): + try: + group._check_owner() + except BaseException as error: # noqa: BLE001 + errors.append(error) + + worker = threading.Thread(target=submit) + worker.start() + worker.join() + assert errors == [] + + current.stream = 45 + with pytest.raises(RuntimeError, match="MAIN stream cannot change"): + group._check_owner() + assert group._execution_stream_handle == 44 + + def _module_constants(path: Path, names: set[str]) -> dict: tree = ast.parse(path.read_text()) constants = { @@ -1648,6 +1702,31 @@ def buffer(ptr): assert to_cute_ptr_cached(buffer(4096)) is not first +def test_executors_of_one_engine_share_the_execution_stream(): + # The planner binds the first MAIN stream for good, so the KV-cache + # estimation executor and the final executor must share one stream. + tree = ast.parse((_TORCH_ROOT / "pyexecutor/py_executor_creator.py").read_text()) + create = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.FunctionDef) and node.name == "_create_py_executor" + ) + calls = [ + node + for node in ast.walk(create) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "create_py_executor_instance" + ] + assert len(calls) >= 2 + for call in calls: + stream = next( + (keyword.value for keyword in call.keywords if keyword.arg == "execution_stream"), + None, + ) + assert isinstance(stream, ast.Name) and stream.id == "execution_stream" + + def test_route_wait_occurs_after_independent_staging(): source = _TORCH_ROOT / "moe/fused_moe/mega_moe/mega_moe_cute_dsl.py" impl = _extract_class( @@ -1791,90 +1870,6 @@ def test_megamoe_runner_identity_is_built_once(): assert builds == [1] -def test_warmup_owner_handoff_uses_deduplicated_manager_public_api(): - events = [] - main = SimpleNamespace(cuda_stream=11, priority=0, device=0) - torch = SimpleNamespace( - cuda=SimpleNamespace( - synchronize=lambda device: events.append(("synchronize", device)), - ) - ) - executor_cls = _extract_class( - _TORCH_ROOT / "pyexecutor/py_executor.py", - "PyExecutor", - {"_handoff_rebalance_warmup_owners"}, - {"torch": torch}, - ) - - class Manager: - def __init__(self, name): - self.name = name - - def handoff_submission_owner(self, stream_handle): - events.append(("handoff", self.name, stream_handle)) - - target = Manager("target") - draft = Manager("draft") - executor = executor_cls() - executor.execution_stream = main - executor.model_engine = SimpleNamespace(moe_load_balancer=target) - executor.draft_model_engine = SimpleNamespace(moe_load_balancer=draft) - - executor._handoff_rebalance_warmup_owners() - assert events == [ - ("synchronize", 0), - ("handoff", "target", 11), - ("handoff", "draft", 11), - ] - - events.clear() - executor.draft_model_engine = SimpleNamespace(moe_load_balancer=target) - executor._handoff_rebalance_warmup_owners() - assert events == [("synchronize", 0), ("handoff", "target", 11)] - - events.clear() - executor.model_engine = SimpleNamespace(moe_load_balancer=None) - executor.draft_model_engine = None - executor._handoff_rebalance_warmup_owners() - assert events == [] - - -def test_quiesced_owner_release_uses_deduplicated_manager_public_api(): - events = [] - main = SimpleNamespace(cuda_stream=11, priority=0, device=0) - executor_cls = _extract_class( - _TORCH_ROOT / "pyexecutor/py_executor.py", - "PyExecutor", - {"_release_quiesced_rebalance_submission_owners"}, - {}, - ) - - class Manager: - def __init__(self, name): - self.name = name - - def release_quiesced_submission_owner(self, stream_handle): - events.append(("release", self.name, stream_handle)) - - target = Manager("target") - draft = Manager("draft") - executor = executor_cls() - executor.execution_stream = main - executor.model_engine = SimpleNamespace(moe_load_balancer=target) - executor.draft_model_engine = SimpleNamespace(moe_load_balancer=draft) - - executor._release_quiesced_rebalance_submission_owners() - assert events == [ - ("release", "target", 11), - ("release", "draft", 11), - ] - - events.clear() - executor.draft_model_engine = SimpleNamespace(moe_load_balancer=target) - executor._release_quiesced_rebalance_submission_owners() - assert events == [("release", "target", 11)] - - def _autotune_helpers(torch): source = _TORCH_ROOT / "moe/custom_ops/cute_dsl_megamoe_custom_op.py" functions = { @@ -2097,6 +2092,301 @@ def close_collectively(self, *, local_safe): ] +def test_post_moe_hook_releases_a_replica_plan_every_chunk(): + moe_cls = _extract_class( + _TORCH_ROOT / "moe/fused_moe/interface.py", + "MoE", + {"_load_balancer_start_set_cpu_stage", "_load_balancer_uses_replica_plan"}, + {}, + ) + moe_cls._using_dynamic_load_balancer = _extract_function( + _TORCH_ROOT / "moe/fused_moe/impl_blocks.py", "_using_dynamic_load_balancer", {} + ) + started = [] + + class Layer: + def __init__(self, *, produces_replica_plan, dynamic_routing): + self.produces_replica_plan = produces_replica_plan + self.dynamic_routing = dynamic_routing + + def is_dynamic_routing(self): + return self.dynamic_routing + + def start_set_cpu_stage(self): + started.append(self) + + cases = [ + # A per-layer dynamic layer disables standard host migration. + (Layer(produces_replica_plan=True, dynamic_routing=False), [True, True]), + (Layer(produces_replica_plan=False, dynamic_routing=True), [False, True]), + (Layer(produces_replica_plan=False, dynamic_routing=False), [False, False]), + (None, [False, False]), + ] + for layer, expected in cases: + moe = moe_cls() + moe.layer_load_balancer = layer + observed = [] + for is_last_call in (False, True): + started.clear() + moe._load_balancer_start_set_cpu_stage(is_last_call) + observed.append(started == [layer]) + assert observed == expected + + +def test_iteration_context_aborts_the_iteration_of_a_failed_forward(): + source = _TORCH_ROOT / "moe/fused_moe/moe_load_balancer.py" + context_cls = _extract_class(source, "MoeLoadBalancerIterContext", None, {}) + base_cls = _extract_class(source, "BaseMoeLoadBalancer", {"abort_iter"}, {}) + events = [] + + class Manager: + static = False + + def is_static_routing(self): + return self.static + + def set_iter_info(self, enable_statistic, enable_update_weights): + events.append("set_iter_info") + + def start_iter(self): + events.append("start_iter") + + def end_iter(self): + events.append("end_iter") + + def abort_iter(self, error): + events.append(("abort_iter", error)) + + manager = Manager() + with context_cls(manager): + pass + assert events == ["set_iter_info", "start_iter", "end_iter"] + + events.clear() + error = RuntimeError("forward failed") + with pytest.raises(RuntimeError, match="forward failed") as caught: + with context_cls(manager): + raise error + assert caught.value is error + assert events == ["set_iter_info", "start_iter", ("abort_iter", error)] + + # Static placement runs no iteration hook. + events.clear() + manager.static = True + with pytest.raises(RuntimeError, match="forward failed"): + with context_cls(manager): + raise error + assert events == [] + + # By default an aborted iteration simply ends, as standard EPLB always did. + class Standard(base_cls): + def end_iter(self): + events.append("end_iter") + + Standard().abort_iter(error) + assert events == ["end_iter"] + + +def test_per_layer_dynamic_abort_releases_pending_plans_and_keeps_the_error(): + namespace = {} + _extract_function( + _TORCH_ROOT / "moe/fused_moe/per_layer_dynamic_eplb/layer.py", "_append_note", namespace + ) + manager_cls = _extract_class( + _TORCH_ROOT / "moe/fused_moe/per_layer_dynamic_eplb/manager.py", + "PerLayerDynamicMoeLoadBalancer", + {"abort_iter", "shutdown"}, + namespace, + ) + released = [] + + class Layer: + def __init__(self, layer_idx, plan=None, cleanup_error=None): + self.layer_idx = layer_idx + self.replica_plan = plan + self.cleanup_error = cleanup_error + + def finish_replica_plan(self): + released.append(self.layer_idx) + if self.cleanup_error is not None: + raise self.cleanup_error + self.replica_plan = None + + manager = manager_cls() + manager.in_iter = True + manager.iter_id = 3 + manager.single_layer_load_balancers = [ + Layer(0), + Layer(1, plan="first", cleanup_error=RuntimeError("route wait failed")), + Layer(2, plan="second"), + ] + error = ValueError("forward failed") + + manager.abort_iter(error) + + # Every pending plan is released; a failed release only annotates the + # forward's error, which the iteration context then re-raises. + assert released == [1, 2] + assert [layer.replica_plan for layer in manager.single_layer_load_balancers] == [ + None, + "first", + None, + ] + assert len(error.__notes__) == 1 + assert "layer 1" in error.__notes__[0] + assert "route wait failed" in error.__notes__[0] + assert (manager.in_iter, manager.iter_id) == (False, 4) + + # The plan left pending makes this rank unsafe for collective teardown. + states = [] + + class Comm: + def allgather(self, state): + states.append(state) + return [state, {"local_safe": True, "pool_count": 0}] + + manager.is_shutdown = False + manager._ep_comm = Comm() + manager._rebalance_shared_slot_pool_order = [] + assert manager.shutdown(local_safe=True) is False + assert states == [{"local_safe": False, "pool_count": 0}] + + # The iteration ends even if the release loop itself fails. + class Broken: + @property + def replica_plan(self): + raise RuntimeError("layer state unavailable") + + manager.in_iter = True + manager.single_layer_load_balancers = [Broken()] + with pytest.raises(RuntimeError, match="layer state unavailable"): + manager.abort_iter(ValueError("forward failed")) + assert (manager.in_iter, manager.iter_id) == (False, 5) + + +class _ChunkRoutes: + dtype = "int32" + + def __init__(self, name): + self.name = name + + +@pytest.mark.parametrize("mode", ["per_layer_dynamic", "standard", "off"]) +def test_fused_comm_chunk_routes_once_after_recording_the_planner_input(mode): + source = _TORCH_ROOT / "moe/fused_moe/moe_scheduler.py" + trace = [] + + class ExpertStatistic: + @staticmethod + def set_layer(layer_idx): + assert layer_idx == 7 + + @staticmethod + def maybe_add_info(num_slots, routes): + trace.append(("statistic", routes.name)) + + def maybe_collect_or_replay_slots(num_slots, routes): + trace.append(("calibrator", routes.name)) + return _ChunkRoutes("replayed " + routes.name) + + namespace = { + "ExpertStatistic": ExpertStatistic, + "get_calibrator": lambda: SimpleNamespace( + maybe_collect_or_replay_slots=maybe_collect_or_replay_slots + ), + "get_active_route_capture": lambda: None, + "torch": SimpleNamespace(int32="int32", float32="float32"), + "Fp4QuantizedTensor": type("Fp4QuantizedTensor", (), {}), + "MoERunContext": lambda **kwargs: SimpleNamespace(**kwargs), + } + _extract_function(source, "_record_routes", namespace) + _extract_function(source, "_leading_rows", namespace) + scheduler_cls = _extract_class(source, "FusedCommMoEScheduler", {"_forward_chunk"}, namespace) + + plan = object() + layer = None if mode == "off" else SimpleNamespace(produces_replica_plan=mode != "standard") + + def route(routes, use_dp): + trace.append(("route", routes.name)) + return _ChunkRoutes("physical" if mode == "per_layer_dynamic" else "slots") + + def run_moe(ctx): + trace.append(("run_moe", ctx.token_selected_experts.name, ctx.replica_plan)) + return "out" + + moe = SimpleNamespace( + apply_router_weight_on_input=False, + mapping=SimpleNamespace(moe_ep_rank=0), + layer_idx=7, + num_slots=24, + use_dp=True, + layer_load_balancer=layer, + routing_method=SimpleNamespace( + apply=lambda logits, input_ids: ( + _ChunkRoutes("logical"), + SimpleNamespace(dtype="float32"), + ) + ), + _maybe_get_perfect_router_expert_ids=lambda routes: routes, + _load_balancer_start_wait_gpu_stage=lambda is_first_call: None, + _load_balancer_done_wait_gpu_stage=lambda is_first_call: None, + _load_balancer_update_statistic=lambda routes, *args, **kwargs: trace.append( + ("update_statistic", routes.name) + ), + _load_balancer_uses_replica_plan=lambda: mode == "per_layer_dynamic", + _load_balancer_route=route, + _load_balancer_replica_plan=lambda: plan if mode == "per_layer_dynamic" else None, + _load_balancer_start_set_cpu_stage=lambda is_last_call: trace.append("start_set_cpu"), + _load_balancer_done_set_cpu_stage=lambda is_last_call: trace.append("done_set_cpu"), + backend=SimpleNamespace( + quantize_input=lambda x, **kwargs: trace.append(("quantize", kwargs)) or ("x", "sf"), + set_adaptive_launch_tokens=lambda tokens: trace.append(("bucket", tokens)), + run_moe=run_moe, + ), + ) + scheduler = scheduler_cls() + scheduler.moe = moe + x = SimpleNamespace(shape=(4, 16)) + + out = scheduler._forward_chunk( + x, + SimpleNamespace(shape=(4, 24)), + output_dtype="bf16", + all_rank_num_tokens=[4, 6], + do_finalize=True, + ) + + assert out == "out" + tail = [("quantize", {}), ("bucket", 6)] + if mode == "per_layer_dynamic": + # The planner consumes the logical routes the calibrator records or replays. + expected = [ + ("update_statistic", "logical"), + ("statistic", "logical"), + ("calibrator", "logical"), + ("route", "replayed logical"), + *tail, + ("run_moe", "physical", plan), + ] + elif mode == "standard": + expected = [ + ("update_statistic", "logical"), + ("route", "logical"), + ("statistic", "slots"), + ("calibrator", "slots"), + *tail, + ("run_moe", "replayed slots", None), + ] + else: + expected = [ + ("statistic", "logical"), + ("calibrator", "logical"), + *tail, + ("run_moe", "replayed logical", None), + ] + assert trace == expected + ["start_set_cpu", "done_set_cpu"] + + def test_base_worker_finalizes_rank_follower_pyexecutor(): worker_cls = _extract_class( _ROOT / "tensorrt_llm/executor/base_worker.py", diff --git a/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py b/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py index eb0984380c04..5c0f28b30284 100644 --- a/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py +++ b/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py @@ -16,9 +16,6 @@ maybe_create_moe_load_balancer, ) from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb import PerLayerDynamicMoeLoadBalancer -from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb.planner import ( - RebalanceSlotSchedulerGroup, -) pytestmark = pytest.mark.cpu_only @@ -79,21 +76,16 @@ class _Group: helper_slots = 2 def __init__(self): - self.has_submission_owner = False self.wait_calls = 0 self.finish_calls = [] - self.handoff_calls = [] - self.release_calls = [] self.events = [] def plan(self, routes, *, defer_wait): assert defer_wait - self.has_submission_owner = True self.events.append("plan") return routes + 4, 7 def plan_schedule(self, routes): - self.has_submission_owner = True self.events.append("schedule") return routes @@ -111,14 +103,6 @@ def wait_for_routes(self): def finish(self, completion_event=None): self.finish_calls.append(completion_event) - def release_warmup_owner_after_sync(self, stream_handle): - self.handoff_calls.append(stream_handle) - self.has_submission_owner = False - - def release_quiesced_submission_owner(self, stream_handle): - self.release_calls.append(stream_handle) - self.has_submission_owner = False - def close(self): pass @@ -163,6 +147,20 @@ def test_layer_produces_and_finishes_replica_plan(): assert layer.replica_plan is None +def test_post_moe_hook_releases_the_chunk_plan(): + manager = _manager() + layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) + group = _Group() + layer.bind_runtime(arena=_Arena(), scheduler_group=group) + + layer.route(torch.tensor([[0, 1]], dtype=torch.int32)) + layer.start_set_cpu_stage() + + assert group.wait_calls == 1 + assert group.finish_calls == [None] + assert layer.replica_plan is None + + def test_overlap_fn_runs_between_schedule_and_plan_finish(): manager = _manager() layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) @@ -315,7 +313,24 @@ def test_iteration_context_reuses_common_lifecycle(): assert get_moe_load_balancer() is None -def test_public_warmup_handoff_and_model_owned_shutdown(): +def test_iteration_context_releases_the_plan_of_a_failed_forward(): + manager = _manager() + layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) + group = _Group() + layer.bind_runtime(arena=_Arena(), scheduler_group=group) + + with pytest.raises(RuntimeError, match="MoE kernel failed"): + with MoeLoadBalancerIterContext(manager, True, True): + layer.route(torch.tensor([[0, 1]], dtype=torch.int32)) + raise RuntimeError("MoE kernel failed") + + assert group.finish_calls == [None] + assert layer.replica_plan is None + assert not manager.in_iter + assert manager.iter_id == 1 + + +def test_model_owned_shutdown(): manager = _manager() layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) arena = _Arena() @@ -324,48 +339,13 @@ def test_public_warmup_handoff_and_model_owned_shutdown(): arena.provider.pool = pool manager._rebalance_shared_slot_pools[("test",)] = pool manager._rebalance_shared_slot_pool_order.append(pool) - group = _Group() - group.has_submission_owner = True - layer.bind_runtime(arena=arena, scheduler_group=group) + layer.bind_runtime(arena=arena, scheduler_group=_Group()) - assert manager.has_submission_owner - manager.handoff_submission_owner(1234) - assert group.handoff_calls == [1234] - manager.release_quiesced_submission_owner(1234) - assert group.release_calls == [1234] assert manager.shutdown() pool.close_collectively.assert_called_once_with(local_safe=True) assert manager.is_shutdown -def _scheduler_group_for_owner_release(*, finished=True): - group = RebalanceSlotSchedulerGroup.__new__(RebalanceSlotSchedulerGroup) - group._closed = False - group._plan_part = None - group._generation = 3 - group._finished_generation = 3 if finished else 2 - group.plan_calls = 3 - group._owner_thread_id = 99 - group._execution_stream_handle = 1234 - return group - - -def test_quiesced_owner_release_clears_worker_and_stream_binding(): - group = _scheduler_group_for_owner_release() - - group.release_quiesced_submission_owner(1234) - - assert group._owner_thread_id is None - assert group._execution_stream_handle is None - - -def test_quiesced_owner_release_rejects_unfinished_generation(): - group = _scheduler_group_for_owner_release(finished=False) - - with pytest.raises(RuntimeError, match="all plans to be finished"): - group.release_quiesced_submission_owner(1234) - - def test_ep_communicator_is_model_owned(monkeypatch): from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb import arena as arena_module From 39a91740dfe71383b835d563c71501b91ce59312 Mon Sep 17 00:00:00 2001 From: jintaop Date: Fri, 23 Oct 2026 01:53:22 -0700 Subject: [PATCH 36/59] [None][refactor] Keep per-layer dynamic EPLB state out of the MegaMoE CuteDSL backend The MegaMoE CuteDSL backend no longer reads the load-balancer configuration or holds per-layer dynamic EPLB state. It keeps only what it needs to consume a replica plan: - the supports_per_layer_dynamic_eplb capability bit; - compute_slot_count: H plus the helper slots of the per-layer dynamic EPLB layer handle (H without one), used only by the live planes, the symmetric workspace and the launch geometry; - quantize_input(reserved_sms=...), which the scheduler fills from the replica plan, plus the plan's launch arguments and route wait. Its create_weights and cache_derived_state match main again. The DWDP check is dropped as well, since MoE initialization already rejects any load balancer under DWDP. The quant method owns the seven live planes from its plane spec. Its create_weights allocates the arena through the layer handle before it registers any Parameter and binds the planes to it at the end; its cache_derived_state checks the planes and has the handle build its HALO-Q/TMA producer. The handle keeps the arena, so allocate_arena is idempotent, bind_backend works from that arena, and live_weight_arena moves from the backend to the handle. The developer guide now describes the package, the lifecycle seams and this split. The per-chunk CUDA work is unchanged. Signed-off-by: jintaop --- .../moe/fused_moe/MOE_DEVELOPER_GUIDE.md | 46 +-- .../fused_moe/mega_moe/mega_moe_cute_dsl.py | 239 +++------------- .../_torch/moe/fused_moe/moe_scheduler.py | 7 +- .../fused_moe/per_layer_dynamic_eplb/layer.py | 47 +-- .../per_layer_dynamic_eplb/planner.py | 45 --- .../_torch/moe/fused_moe/quantization.py | 89 +++++- .../_torch/moe/test_dynamic_eplb_contracts.py | 66 ++++- tests/unittest/_torch/moe/test_moe_backend.py | 270 ++++++++++++++---- .../_torch/moe/test_per_layer_dynamic_eplb.py | 66 +++-- 9 files changed, 508 insertions(+), 367 deletions(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md index 9bab613bc800..b4e43723d913 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -221,6 +221,7 @@ Still on old path (standalone, with embedded communication): | `routing.py` | Routing methods (`TopKRouting`, etc.) | | `moe_comm_timeout_guard.py` | Device-side timeouts of communication kernels, relaxed during warmup — budget resolution and `MoECommTimeoutProxy` registration (see [Communication timeouts](#communication-timeouts)) | | `moe_load_balancer.py` | EPLB implementation | +| `per_layer_dynamic_eplb/` | Per-layer dynamic EPLB, created by `maybe_create_moe_load_balancer`: model-level `manager`, per-layer handle `layer`, live-weight `arena`, and HALO-Q/TMA `planner` | | `moe_op_backend.py` | Op backend registry for TRTLLMGen (flashinfer/trtllm ops) | ### Backends (`fused_moe/`) @@ -876,9 +877,10 @@ set wrong has a test that catches it: copy stream, one CTA per auxiliary SM. Each CTA raises a device route epoch once its routes are written. The main stream may enqueue shared-expert work and input quantization before its first routed-expert use, where a stream - memory wait on that epoch orders it after the routes. A single CPU submitter - owns both streams; the path does not require a CUDA green context or a - steady-state host wait. + memory wait on that epoch orders it after the routes. Each layer binds one + main stream, the executor's execution stream, on which warmup and serving + submit from their own threads; the path does not require a CUDA green + context or a steady-state host wait. - HALO-Q waits for every EP rank inside the kernel and traps once a peer is later than its spin bound. Each launch passes the bound, which `SchedulerSpinBound` takes from the MoE communication timeout guard: the @@ -897,20 +899,32 @@ set wrong has a test that catches it: rollback must release the scheduler, TMA bindings, shared-slot mappings, and fabric allocations collectively. Forward execution must not retain activation tensors after their consumer has been enqueued. -- Per-layer dynamic EPLB reuses the EPLB seams. `maybe_create_moe_load_balancer` - builds `PerLayerDynamicMoeLoadBalancer`, which shares the - `BaseMoeLoadBalancer` lifecycle with `MoeLoadBalancer` (context entry, - `set_iter_info`, `start_iter`/`end_iter`, collective `shutdown`); warmup - reaches it only through `set_iter_info`. `add_layer` returns a - `PerLayerDynamicMoeLayerLoadBalancer` with `produces_replica_plan = True`: - `_load_balancer_route` plans, the backend reads `MoERunContext.replica_plan`, - and the scheduler finishes the plan after each chunk's `run_moe`. Model code - passes independent main-stream work such as shared experts through +- Per-layer dynamic EPLB lives in `per_layer_dynamic_eplb/` (`manager`, + `layer`, `arena`, `planner`) and reuses the EPLB seams. + `maybe_create_moe_load_balancer` builds `PerLayerDynamicMoeLoadBalancer`, + which shares the `BaseMoeLoadBalancer` lifecycle with `MoeLoadBalancer`, and + `add_layer` returns a `PerLayerDynamicMoeLayerLoadBalancer` with + `produces_replica_plan = True`. In every chunk `_load_balancer_route` plans, + the backend reads `MoERunContext.replica_plan`, and the post-`run_moe` hook + releases the plan. The planner consumes logical routes, so the expert + statistics and the layer-wise calibrator record (or replay) them before the + route; standard EPLB records the slots after it. When the forward raises, + `MoeLoadBalancerIterContext` calls `abort_iter`, which releases the pending + plans. Plans run during warmup too, so there is no warmup switch; only the + HALO-Q spin bound changes, as described above. Model code passes independent + main-stream work such as shared experts through `MoE.load_balancer_overlap(fn)`; the route runs it once, after submitting the - copy stream. The backend allocates its live-weight arena through the layer - handle's `allocate_arena` and exposes `live_weight_arena`, - `ep_process_group`, and `max_tokens_per_rank_capacity` to `bind_backend` - instead of private attributes. + copy stream. +- The quant method owns the seven live kernel planes. Its `create_weights` + allocates them from its plane spec through the layer handle's + `allocate_arena` and binds them; its `cache_derived_state` checks them and + calls the handle's `bind_backend`, which builds the HALO-Q/TMA producer. The + backend only declares `supports_per_layer_dynamic_eplb`, uses + `compute_slot_count` (H + S, with S from the layer handle) for the planes, + symmetric workspace, and launch geometry while `expert_size_per_partition` + stays H, and consumes the replica plan, including the SMs its input + quantization leaves to the copy kernel. `bind_backend` reads the backend's + `ep_process_group` and `max_tokens_per_rank_capacity`. - DeepSeek-V4 selects the fused shared FC12 kernel when the routed backend declares per-layer dynamic support, the final projection quantization and hardware (SM100 family) support that kernel, and the model runs eagerly diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py index 550cbe11306d..1329c405c95b 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py @@ -68,6 +68,7 @@ from __future__ import annotations +import functools import os import socket import weakref @@ -140,8 +141,6 @@ class MegaMoeCuteDslUnavailable(RuntimeError): _MEGAMOE_PRIME_LADDER = os.environ.get("MEGAMOE_PRIME_LADDER", "1") == "1" _MEGAMOE_PRIMED_LADDERS: set = set() -# Distinguish an absent arm attribute from an explicit False when restoring it. - def is_megamoe_cute_dsl_runtime_available() -> Tuple[bool, Optional[str]]: """Return whether the CUDA 13 Cutlass DSL runtime exposes all symbols the @@ -287,6 +286,17 @@ def _as_fp8_sf(t: torch.Tensor) -> torch.Tensor: return t if t.dtype == torch.float8_e4m3fn else t.view(torch.float8_e4m3fn) +@functools.cache +def _fp4_quantize_has_sm_budget() -> bool: + if hasattr(torch.ops.trtllm.fp4_quantize, "sm_budget"): + return True + logger.warning( + "fp4_quantize.sm_budget is unavailable; per-layer dynamic EPLB " + "quantization will use the selected backend without SM control" + ) + return False + + # --------------------------------------------------------------------------- # Weight view passed to ``run_moe`` # --------------------------------------------------------------------------- @@ -434,16 +444,6 @@ class TrtllmCutedslMegaMoeNvfp4Impl(MoEImplBase): # ``CombineFormat.parse`` in the kernel package's token_comm.py. _SUPPORTED_COMBINE_FORMATS = frozenset({"bf16", "32e4m3xe8m0", "16e2m1xbf16"}) - # Disabled defaults also cover construction paths before weight allocation. - _rebalance_slots_active: int = 0 - _quantize_reserved_sms: int = 0 - # Resident count H is recorded at weight creation and never widened. - _rebalance_home_experts: Optional[int] = None - # Active per-layer dynamic variant of ``model_config.moe_load_balancer``. - _per_layer_dynamic_eplb_config = None - # Allocate the arena once; the copy endpoint retains its plane identities. - _rebalance_arena = None - # ------------------------------------------------------------------ # Capability gating # ------------------------------------------------------------------ @@ -556,23 +556,6 @@ def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: # ------------------------------------------------------------------ # Init # ------------------------------------------------------------------ - @staticmethod - def _reject_rebalance_dwdp_overlap(mapping, helper_slots: int, layer_idx=None) -> None: - """Reject simultaneous helper-slot rebalance and DWDP before allocation. - - DWDP makes every rank own all experts and rewrites the slot axis after - weight creation, invalidating the H + S arena geometry. - """ - if int(helper_slots) <= 0: - return - if not bool(getattr(mapping, "dwdp_enabled", False)): - return - raise ValueError( - f"MoE rebalance S={helper_slots} at layer {layer_idx} cannot run with " - f"DWDP size {getattr(mapping, 'dwdp_size', '?')}: DWDP rewrites the " - "H + S weight axis. Disable rebalance or use dwdp_size=1." - ) - def __init__( self, *, @@ -644,61 +627,6 @@ def __init__( # gates collective tuning behavior). self.tactic_autotune = os.environ.get("MEGAMOE_TACTIC_AUTOTUNE", "0") == "1" - # Per-layer dynamic EPLB uses the standard load-balancer configuration. - # P is the global physical slot count; the fused kernel sees - # M = P / EP slots per rank and S = M - H helper rows. - load_balancer = getattr(model_config, "moe_load_balancer", None) - self._per_layer_dynamic_eplb_config = ( - load_balancer if getattr(load_balancer, "mode", None) == "per_layer_dynamic" else None - ) - self._rebalance_slots_active = 0 - if self._per_layer_dynamic_eplb_config is not None: - ep_size = int(self.ep_size) - global_slots = int(self._per_layer_dynamic_eplb_config.num_slots) - if ep_size <= 0 or num_experts <= 0 or num_experts % ep_size != 0: - raise ValueError( - "per-layer dynamic EPLB requires a positive expert count divisible " - f"by EP; got E={num_experts}, EP={ep_size}." - ) - if global_slots % ep_size != 0: - raise ValueError( - "per-layer dynamic EPLB requires num_slots divisible by EP; " - f"got P={global_slots}, EP={ep_size}." - ) - resident_experts = num_experts // ep_size - compute_slots = global_slots // ep_size - self._rebalance_slots_active = compute_slots - resident_experts - if self._rebalance_slots_active <= 0: - raise ValueError( - "per-layer dynamic EPLB requires physical capacity larger than " - f"resident capacity; got P={global_slots}, E={num_experts}." - ) - - self._quantize_reserved_sms = 0 - self._rebalance_reserved_sms = 0 - if self._rebalance_slots_active > 0: - self._rebalance_reserved_sms = int(self._per_layer_dynamic_eplb_config.auxiliary_sms) - if hasattr(torch.ops.trtllm.fp4_quantize, "sm_budget"): - self._quantize_reserved_sms = self._rebalance_reserved_sms - else: - logger.warning_once( - "fp4_quantize.sm_budget is unavailable; per-layer dynamic EPLB " - "quantization will use the selected backend without SM control", - key="dynamic_eplb_quantize_without_sm_budget", - ) - # Reject DWDP before it can rewrite the same slot axis. - self._reject_rebalance_dwdp_overlap( - getattr(self, "mapping", None), - self._rebalance_slots_active, - layer_idx, - ) - if self._rebalance_slots_active > 0: - logger.debug( - f"[MegaMoECuteDsl] prefill rebalance active " - f"(S={self._rebalance_slots_active}); kernel weight planes " - f"use M = H + S while expert_size_per_partition stays H." - ) - # Buffer sizing. MoE layers execute serially per forward; one pool # sized to the worst-case per-rank tokens covers every layer. The # kernel compile takes this as the static ``max_tokens_per_rank``. @@ -787,26 +715,19 @@ def compute_slot_count(self) -> int: """Return the kernel-only slot axis ``M = H + S``. ``expert_size_per_partition`` remains the resident/checkpoint axis H. - Only the seven live kernel planes, symmetric workspace, and launch - geometry use this wider axis. + The per-layer dynamic EPLB layer handle supplies the S helper slots; + without one, M is H. Only the seven live kernel planes, symmetric + workspace, and launch geometry use this wider axis. """ resident_experts = int(self.expert_size_per_partition) - recorded_home = self._rebalance_home_experts - if recorded_home is not None and resident_experts != recorded_home: - raise RuntimeError( - "MegaMoECuteDsl resident expert geometry changed after weight " - f"creation: expected H={recorded_home}, got H={resident_experts}." - ) - return resident_experts + int(self._rebalance_slots_active) + layer_load_balancer = self.layer_load_balancer + if layer_load_balancer is None or not layer_load_balancer.produces_replica_plan: + return resident_experts + return resident_experts + int(layer_load_balancer.helper_slots_per_rank) # ------------------------------------------------------------------ - # Facts the per-layer dynamic EPLB layer handle binds against + # Backend facts the per-layer dynamic EPLB layer handle reads # ------------------------------------------------------------------ - @property - def live_weight_arena(self): - """The H + S arena that the live kernel planes alias, or ``None``.""" - return self._rebalance_arena - @property def ep_process_group(self): """The EP ProcessGroup resolved at construction, or ``None`` on one rank.""" @@ -1069,15 +990,12 @@ def create_weights(self): """Build-time weight + symmetric-buffer allocation. Order: - 0. Record the resident/checkpoint axis H. The backend keeps - ``expert_size_per_partition == H``; ``compute_slot_count`` exposes - ``M = H + S`` only to kernel-plane allocation and launch geometry. 1. Allocate symmetric-memory provider for multi-rank EP (collective rendezvous; MUST run at build time -- not from ``run_moe`` -- because forward time may be inside CUDA graph capture or non-lockstep PP/layer-skip). - 2. Resolve the quantization method and its seven-plane ABI. - 3. Delegate parameter registration to the same quant method. + 2. Resolve quantization method. + 3. Delegate parameter registration to the quant method. 4. Flip ``_weights_created``. The symm provider is shared across MoE layers with the same @@ -1089,62 +1007,13 @@ def create_weights(self): """ if self._weights_created: return - # ConfigurableMoE has synchronized resident geometry H before this call. - # Keep that public/checkpoint axis unchanged. The explicit - # ``compute_slot_count`` property is the only H + S axis. - resident_experts = int(self.expert_size_per_partition) - helper_slots = int(self._rebalance_slots_active) - if self._rebalance_home_experts is None: - self._rebalance_home_experts = resident_experts - elif resident_experts != self._rebalance_home_experts: - raise RuntimeError( - "MegaMoECuteDsl: expert_size_per_partition changed between " - f"create_weights() calls -- expected resident H=" - f"{self._rebalance_home_experts}, found H={resident_experts}." - ) - if helper_slots > 0: - log_fn = logger.info if self.layer_idx in (0, None) else logger.debug - log_fn( - f"[MegaMoECuteDsl] layer={self.layer_idx} MoE rebalance uses " - f"resident H={resident_experts} and helper S={helper_slots}; " - f"only the seven live kernel planes use " - f"M={self.compute_slot_count} rows." - ) - # The quant method is the sole authority for the seven live planes. - if self.quant_method is None: - self.quant_method = self._get_quant_method() - live_weight_plane_spec = self.quant_method.live_weight_plane_spec(self) - # Step 1: build-time symmetric memory allocation (multi-rank only). # Single-rank degenerate uses local CUDA tensors and skips here. if self.ep_size > 1: self._alloc_symm_provider() - # Allocate the H + S arena before the quant method creates the seven - # kernel-plane Parameters. Source/checkpoint Parameters remain H rows. - if self._rebalance_slots_active > 0 and self._rebalance_arena is None: - # The layer handle allocates the hierarchical multicast arena (one - # multicast group per HALO level, rather than one per source and - # destination pair) from model-owned resources. - layer_balancer = getattr(self, "layer_load_balancer", None) - allocate_arena = getattr(layer_balancer, "allocate_arena", None) - if not callable(allocate_arena): - raise RuntimeError( - "per-layer dynamic EPLB requires a layer load balancer that " - "allocates the live-weight arena" - ) - self._rebalance_arena = allocate_arena( - bundle=live_weight_plane_spec, - ep_process_group=self._ep_pg, - mapping=self.mapping, - device=int(self.mapping.local_rank), - layer_idx=self.layer_idx, - ) - # Step 2-3: the same quant method registers every live Parameter. + # Step 2-3: quant method registers all NVFP4 + MegaMoE-format params. + self.quant_method = self._get_quant_method() self.quant_method.create_weights(self) - if self._rebalance_arena is not None: - from ..per_layer_dynamic_eplb.planner import bind_live_weight_planes - - bind_live_weight_planes(self.quant_method, self, self._rebalance_arena) # Step 4. self._weights_created = True @@ -1303,17 +1172,6 @@ def cache_derived_state(self) -> None: if self.quant_method is None: self.create_weights() super().cache_derived_state() - # ConfigurableMoE delegates this hook on both ordinary and staged - # loading paths. The layer manager builds and owns the producer only - # after the final seven plane aliases exist. - if self._rebalance_slots_active > 0: - arena = self._rebalance_arena - arena.assert_identity() - self.quant_method.assert_live_weight_aliases(self, arena) - bind_backend = getattr(self.layer_load_balancer, "bind_backend", None) - if not callable(bind_backend): - raise RuntimeError("per-layer dynamic EPLB layer manager cannot bind its runtime") - bind_backend(self) def pre_reload_weights(self) -> None: raise NotImplementedError( @@ -1327,7 +1185,7 @@ def _build_weight_view(self) -> MegaMoECuteDslWeightView: kernel sees the latest dynamic-EPLB migration outcome (once that path lands; currently the slots are static). """ - _view = MegaMoECuteDslWeightView( + return MegaMoECuteDslWeightView( fc1_weight=self.mega_fc1_weight, fc1_weight_sf=self.mega_fc1_weight_sf, fc2_weight=self.mega_fc2_weight, @@ -1336,7 +1194,6 @@ def _build_weight_view(self) -> MegaMoECuteDslWeightView: fc2_alpha=self.fc2_alpha, fc1_norm_const=self.fc1_norm_const, ) - return _view # ------------------------------------------------------------------ # MoE-contract methods @@ -1346,6 +1203,7 @@ def quantize_input( x: Union[torch.Tensor, "Fp4QuantizedTensor"], *, post_quant_comm: bool = False, + reserved_sms: int = 0, **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor]: """BF16 -> NVFP4 packed activation + plain K-major FP8 SF. @@ -1363,6 +1221,10 @@ def quantize_input( ``fc31_input_scale`` was calibrated on the scaled activation. The kernel itself only ever sees the quantized input. + ``reserved_sms`` is the number of SMs left to the concurrent + HALO-Q/TMA kernel of a replica plan; zero uses the default + quantization overload. + Empty input (``x.shape[0] == 0``) short-circuits to empty NVFP4 + empty SF without launching a quantization kernel, so the ``FusedCommMoEScheduler`` can call ``quantize_input`` uniformly @@ -1390,16 +1252,16 @@ def quantize_input( # admits bf16 activations only, so ``fc31_act_scale``, created at # ``module.dtype``, is bf16 as well. x_bf16 = x_bf16 * fc31_act_scale - # Match the concurrent scheduler/copy reservation. The default overload - # remains unchanged when load balancing is disabled. - if self._quantize_reserved_sms: + # Leave a replica plan's SMs to its concurrent HALO-Q/TMA kernel. The + # default overload remains unchanged without a plan. + if reserved_sms > 0 and _fp4_quantize_has_sm_budget(): x_fp4, x_sf = torch.ops.trtllm.fp4_quantize.sm_budget( x_bf16, self.fc31_input_scale, 16, False, False, - reserved_sms=self._quantize_reserved_sms, + reserved_sms=reserved_sms, ) else: x_fp4, x_sf = torch.ops.trtllm.fp4_quantize( @@ -1446,13 +1308,6 @@ def _replica_plan_launch_args(self, replica_plan): f"S={helper_slots}, generation={generation}, " f"reserved_sms={reserved_sms}." ) - quantize_reserved_sms = int(getattr(self, "_quantize_reserved_sms", 0)) - if quantize_reserved_sms not in (0, reserved_sms): - raise RuntimeError( - "MegaMoECuteDsl replica-plan SM reservation differs from the " - f"quantization reservation: plan={reserved_sms}, " - f"quantize={quantize_reserved_sms}." - ) return ( compute_slots, helper_slots, @@ -1736,7 +1591,6 @@ def _stage_inputs( num_tokens: int, top_k: int, stage_activation: bool, - stage_routes: bool = True, replica_plan=None, ) -> None: """Copy live rows of the user-domain inputs into the kernel's @@ -1764,7 +1618,7 @@ def _stage_inputs( """ max_T = bufs.topk_idx_local.shape[0] last_T = self._last_staged_T.get(max_T) - if stage_routes and last_T is not None and last_T > num_tokens: + if last_T is not None and last_T > num_tokens: bufs.topk_idx_local[num_tokens:last_T].fill_(-1) if num_tokens > 0: if stage_activation: @@ -1773,11 +1627,10 @@ def _stage_inputs( bufs.topk_weights[:num_tokens, :top_k].copy_(topk_weights, non_blocking=True) if num_tokens < max_T: bufs.topk_weights[num_tokens:max_T, :top_k].zero_() - if stage_routes: - if num_tokens > 0: - self._wait_rebalance_routes(replica_plan) - bufs.topk_idx_local[:num_tokens].copy_(topk_idx, non_blocking=True) - self._last_staged_T[max_T] = num_tokens + if num_tokens > 0: + self._wait_rebalance_routes(replica_plan) + bufs.topk_idx_local[:num_tokens].copy_(topk_idx, non_blocking=True) + self._last_staged_T[max_T] = num_tokens def _launch_megamoe_kernel( self, @@ -1878,8 +1731,6 @@ def _deferred_scratch_factory(_self=self, _get=get_megamoe_profiling_scratch): helper_ready_generation, _reserved_sms, ) = plan_args - # Record the actual helper-slot ABI for verification. - self._last_S_to_kernel = _S_active # Quantization, independent staging, output clears, and host setup # are already enqueued. The router inside this op first reads routes. self._wait_rebalance_routes(replica_plan) @@ -2017,7 +1868,7 @@ def _prime_ladder_buckets( weight_view, primed, ): - """Prime each max-token bucket; the caller manages the rebalance arm state.""" + """Prime each max-token bucket.""" for max_T in self._maxt_buckets: if max_T == live_T: continue @@ -2151,7 +2002,6 @@ def _run_moe( and topk_idx.data_ptr() % 16 == 0 and topk_weights.data_ptr() % 16 == 0 ) - borrowed_routes = None if not use_direct_routing_inputs: self._stage_inputs( bufs=bufs, @@ -2162,19 +2012,12 @@ def _run_moe( num_tokens=num_tokens, top_k=top_k, stage_activation=not use_direct_inputs, - stage_routes=borrowed_routes is None, replica_plan=replica_plan, ) return self._launch_megamoe_kernel( activation=x if use_direct_inputs else bufs.activation, activation_sf=x_sf if use_direct_inputs else bufs.activation_sf, - topk_idx=( - borrowed_routes - if borrowed_routes is not None - else topk_idx - if use_direct_routing_inputs - else bufs.topk_idx_local - ), + topk_idx=topk_idx if use_direct_routing_inputs else bufs.topk_idx_local, topk_weights=( topk_weights if use_direct_routing_inputs else bufs.topk_weights[:, :top_k] ), diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py index f755e214b8dc..ac99e08f2308 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py @@ -1596,12 +1596,17 @@ def _forward_chunk( # prepare its workspace. moe_input = x_chunk_real x_sf = None - else: + elif replica_plan is None: # Delegate to ``backend.quantize_input`` so each fused-comm backend # owns its own empty-tensor layout. Both MegaMoEDeepGemm and # MegaMoECuteDsl short-circuit ``x.shape[0] == 0`` inside their # quantize_input contracts. moe_input, x_sf = moe.backend.quantize_input(x_chunk_real) + else: + # Leave the plan's SMs to its concurrent HALO-Q/TMA kernel. + moe_input, x_sf = moe.backend.quantize_input( + x_chunk_real, reserved_sms=replica_plan.reserved_sms + ) # CuteDSL needs the scheduler's rank-identical chunk maximum to select # one adaptive bucket on every EP rank; using a local token count could diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/layer.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/layer.py index 4b5607baff38..a1d645109f06 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/layer.py +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/layer.py @@ -144,10 +144,15 @@ def bind_runtime(self, *, arena: Any, scheduler_group: Any) -> None: """Attach resources built after final weight storage is available.""" if arena is None or scheduler_group is None: raise ValueError("arena and scheduler_group are both required") - if self._arena is not None: + if self._scheduler_group is not None: if arena is self._arena and scheduler_group is self._scheduler_group: return raise RuntimeError(f"per-layer dynamic EPLB layer {self.layer_idx} is already bound") + if self._arena is not None and arena is not self._arena: + raise RuntimeError( + f"per-layer dynamic EPLB layer {self.layer_idx} cannot bind an arena " + "it did not allocate" + ) arena_home = getattr(arena, "home_experts", self.resident_experts_per_rank) arena_helpers = getattr(arena, "helper_slots", self.helper_slots_per_rank) @@ -179,14 +184,21 @@ def allocate_arena( ) -> Any: """Allocate this layer's H + S live-weight arena from model-owned resources. - Collective across the EP group: every rank calls it from the same point - of the backend's ``create_weights``, after the symmetric-buffer - rendezvous and before the kernel-plane Parameters are registered. - ``bundle`` is the quant method's live-weight plane spec. + Collective across the EP group: the quant method calls it from its + ``create_weights`` on every rank, before it registers the kernel-plane + Parameters. ``bundle`` is the quant method's live-weight plane spec. + A later call returns the same arena. """ + if self._arena is not None: + if bundle != self._arena.bundle: + raise RuntimeError( + f"per-layer dynamic EPLB layer {self.layer_idx} already allocated " + "its arena for a different live-weight plane spec" + ) + return self._arena from .arena import allocate_rebalance_arena - return allocate_rebalance_arena( + self._arena = allocate_rebalance_arena( home_experts=self.resident_experts_per_rank, helper_slots=self.helper_slots_per_rank, bundle=bundle, @@ -196,13 +208,21 @@ def allocate_arena( device=int(device), layer_idx=layer_idx, ) + return self._arena + + @property + def live_weight_arena(self) -> Optional[Any]: + """The H + S arena that the live kernel planes alias, or ``None``.""" + return self._arena def bind_backend(self, backend: Any) -> None: """Build and own the producer after the backend's live aliases exist.""" - arena = backend.live_weight_arena + arena = self._arena if arena is None: - raise RuntimeError("per-layer dynamic EPLB backend has no live arena") - if self._arena is not None or self._scheduler_group is not None: + raise RuntimeError( + f"per-layer dynamic EPLB layer {self.layer_idx} has no live-weight arena" + ) + if self._scheduler_group is not None: bound_backend = ( self._bound_backend_ref() if self._bound_backend_ref is not None else None ) @@ -211,15 +231,6 @@ def bind_backend(self, backend: Any) -> None: f"per-layer dynamic EPLB layer {self.layer_idx} is already " "bound to a different backend" ) - if self._arena is not arena: - raise RuntimeError( - f"per-layer dynamic EPLB layer {self.layer_idx} backend " - "changed its rebalance arena after binding" - ) - if self._scheduler_group is None: - raise RuntimeError( - f"per-layer dynamic EPLB layer {self.layer_idx} has an incomplete runtime binding" - ) return backend_ref = weakref.ref(backend) diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py index b2fac04d5051..9e0205adc7ba 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py @@ -489,48 +489,3 @@ def build_rebalance_slot_scheduler_group( spin_bound=manager.get_scheduler_spin_bound(), layer_idx=layer_idx, ) - - -def bind_live_weight_planes(quant_method: Any, module: Any, arena: Any) -> None: - """Alias MegaMoE's seven live parameters onto its rebalance arena.""" - expected_bundle = quant_method.live_weight_plane_spec(module) - if expected_bundle != arena.bundle: - raise RuntimeError("MoE rebalance arena does not match the quant-method plane spec") - expected_names = tuple(plane.name for plane in expected_bundle.planes) - if tuple(arena.plane_names) != expected_names: - raise RuntimeError("MoE rebalance arena plane order is not canonical") - rebound = False - for name, arena_view, alias in zip( - arena.plane_names, arena.local_plane_views, arena.tekit_alias_views - ): - old = getattr(module, name, None) - if old is None: - raise RuntimeError(f"MoE rebalance live weight plane {name!r} was not registered") - if ( - tuple(alias.shape) != tuple(old.shape) - or tuple(alias.stride()) != tuple(old.stride()) - or alias.dtype != old.dtype - ): - raise RuntimeError(f"MoE rebalance arena alias for {name!r} has an incompatible layout") - if not alias.is_cuda or int(alias.data_ptr()) != int(arena_view.data_ptr()): - raise RuntimeError(f"MoE rebalance arena alias for {name!r} is not its allocator view") - if old.is_meta: - setattr( - module, - name, - torch.nn.Parameter(alias, requires_grad=old.requires_grad), - ) - rebound = True - else: - try: - old.data = alias - except RuntimeError as error: - raise RuntimeError( - f"MoE rebalance failed to bind live plane {name!r}: {error}" - ) from error - if rebound: - quant_method.setup_quant_scales(module) - norm_const = getattr(module, "fc1_norm_const", None) - if norm_const is not None and arena.home_experts < norm_const.data.shape[0]: - norm_const.data[int(arena.home_experts) :].fill_(1.0) - quant_method.assert_live_weight_aliases(module, arena) diff --git a/tensorrt_llm/_torch/moe/fused_moe/quantization.py b/tensorrt_llm/_torch/moe/fused_moe/quantization.py index 601de04bead5..1607bd2af75f 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/moe/fused_moe/quantization.py @@ -4381,13 +4381,67 @@ def assert_live_weight_aliases(self, module: torch.nn.Module, f"MoE rebalance live weight plane {plane.name!r} changed storage after binding" ) + def _bind_live_weight_planes(self, module: torch.nn.Module, arena) -> None: + """Alias MegaMoE's seven live parameters onto its rebalance arena.""" + expected_bundle = self.live_weight_plane_spec(module) + if expected_bundle != arena.bundle: + raise RuntimeError( + "MoE rebalance arena does not match the quant-method plane spec" + ) + expected_names = tuple(plane.name for plane in expected_bundle.planes) + if tuple(arena.plane_names) != expected_names: + raise RuntimeError( + "MoE rebalance arena plane order is not canonical") + rebound = False + for name, arena_view, alias in zip(arena.plane_names, + arena.local_plane_views, + arena.tekit_alias_views): + old = getattr(module, name, None) + if old is None: + raise RuntimeError( + f"MoE rebalance live weight plane {name!r} was not registered" + ) + if (tuple(alias.shape) != tuple(old.shape) + or tuple(alias.stride()) != tuple(old.stride()) + or alias.dtype != old.dtype): + raise RuntimeError( + f"MoE rebalance arena alias for {name!r} has an incompatible layout" + ) + if (not alias.is_cuda + or int(alias.data_ptr()) != int(arena_view.data_ptr())): + raise RuntimeError( + f"MoE rebalance arena alias for {name!r} is not its allocator view" + ) + if old.is_meta: + setattr( + module, + name, + torch.nn.Parameter(alias, requires_grad=old.requires_grad), + ) + rebound = True + else: + try: + old.data = alias + except RuntimeError as error: + raise RuntimeError( + f"MoE rebalance failed to bind live plane {name!r}: {error}" + ) from error + if rebound: + self.setup_quant_scales(module) + norm_const = getattr(module, "fc1_norm_const", None) + if (norm_const is not None + and arena.home_experts < norm_const.data.shape[0]): + norm_const.data[int(arena.home_experts):].fill_(1.0) + self.assert_live_weight_aliases(module, arena) + # ----------------------------------------------------------------- # create_weights: register MegaMoE-format parameters in addition to # the grandparent's standard NVFP4 parameters. # ----------------------------------------------------------------- def create_weights(self, module: torch.nn.Module): - # Resolve the same descriptor used by build-time arena allocation before - # the parent materializes any source or derived Parameters. + # Resolve the seven-plane descriptor before the parent materializes any + # source or derived Parameters; a per-layer dynamic EPLB arena is + # allocated from it first. bundle = self.live_weight_plane_spec(module) # Fail fast: the derived-only EPLB migration path registers no bias # staging. @@ -4395,6 +4449,19 @@ def create_weights(self, module: torch.nn.Module): raise NotImplementedError( "NVFP4MegaMoECuteDslMethod does not support expert bias " "together with dynamic-EPLB shared weight loading.") + layer_load_balancer = getattr(module, "layer_load_balancer", None) + arena = None + if (layer_load_balancer is not None + and layer_load_balancer.produces_replica_plan): + # Collective over EP: every rank allocates its H + S live-weight + # arena here. + arena = layer_load_balancer.allocate_arena( + bundle=bundle, + ep_process_group=module.ep_process_group, + mapping=module.mapping, + device=int(module.mapping.local_rank), + layer_idx=module.layer_idx, + ) weight_vec_size = torch.iinfo(self.weight_dtype).bits // 4 self.block_scales_vec_size = torch.iinfo( @@ -4492,6 +4559,24 @@ def storage_parameter(name: str, factory=torch.empty) -> nn.Parameter: # Alpha Parameters may have been resized or rebound after the parent # built its view; refresh it against the final seven-plane identities. self.setup_quant_scales(module) + if arena is not None: + self._bind_live_weight_planes(module, arena) + + def cache_derived_state(self, module: torch.nn.Module) -> None: + super().cache_derived_state(module) + layer_load_balancer = getattr(module, "layer_load_balancer", None) + if (layer_load_balancer is None + or not layer_load_balancer.produces_replica_plan): + return + # Loading and transformation must leave the seven planes on the arena; + # the layer handle then builds its HALO-Q/TMA producer against them. + arena = layer_load_balancer.live_weight_arena + if arena is None: + raise RuntimeError( + "per-layer dynamic EPLB layer has no live-weight arena") + arena.assert_identity() + self.assert_live_weight_aliases(module, arena) + layer_load_balancer.bind_backend(module) def _materialize_source_params(self, module: torch.nn.Module): """Rematerialize this module's streamed source params (full shape) diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index 8f3e123389ed..b16ed958c30b 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -1736,7 +1736,7 @@ def test_route_wait_occurs_after_independent_staging(): {"torch": SimpleNamespace(uint8="uint8")}, ) - def run(stage_routes): + def run(): trace = [] owner = impl() owner._last_staged_T = {8: 8} @@ -1756,20 +1756,64 @@ def run(stage_routes): num_tokens=4, top_k=6, stage_activation=True, - stage_routes=stage_routes, ) return trace, owner - trace, owner = run(True) + trace, owner = run() assert trace.index("copy:activation") < trace.index("wait") assert trace.index("copy:topk_weights") < trace.index("wait") assert trace.index("wait") < trace.index("copy:topk_idx_local") assert owner._last_staged_T[8] == 4 - trace, owner = run(False) - assert "wait" not in trace - assert not any("topk_idx_local" in item for item in trace) - assert owner._last_staged_T[8] == 8 + +@pytest.mark.parametrize("has_sm_budget", [True, False]) +def test_megamoe_quantize_reserves_sms_only_for_a_replica_plan(has_sm_budget): + source = _TORCH_ROOT / "moe/fused_moe/mega_moe/mega_moe_cute_dsl.py" + calls = [] + + class Activation: + def __init__(self, rows, cols): + self.shape = (rows, cols) + + def to(self, dtype): + return self + + def contiguous(self): + return self + + def view(self, *shape): + return shape + + def fp4_quantize(*args): + calls.append(("default", *args[1:])) + return "x_fp4", Activation(3, 4) + + def sm_budget(*args, reserved_sms): + calls.append(("sm_budget", *args[1:], reserved_sms)) + return "x_fp4", Activation(3, 4) + + if has_sm_budget: + fp4_quantize.sm_budget = sm_budget + namespace = { + "torch": SimpleNamespace( + bfloat16="bf16", ops=SimpleNamespace(trtllm=SimpleNamespace(fp4_quantize=fp4_quantize)) + ), + "logger": SimpleNamespace(warning=lambda message: None), + "Fp4QuantizedTensor": type("Fp4QuantizedTensor", (), {}), + "megamoe_activation_sf_bytes_per_row": lambda hidden: hidden // 16, + "ceil_div": lambda a, b: -(-a // b), + } + _extract_function(source, "_fp4_quantize_has_sm_budget", namespace) + impl = _extract_class(source, "TrtllmCutedslMegaMoeNvfp4Impl", {"quantize_input"}, namespace) + owner = impl() + owner.fc31_input_scale = "scale" + + owner.quantize_input(Activation(3, 64)) + owner.quantize_input(Activation(3, 64), reserved_sms=4) + + default = ("default", "scale", 16, False, False) + budget = ("sm_budget", "scale", 16, False, False, 4) if has_sm_budget else default + assert calls == [default, budget] def test_megamoe_launch_buffers_are_built_once_per_bucket_and_mode(): @@ -2303,7 +2347,7 @@ def maybe_collect_or_replay_slots(num_slots, routes): _extract_function(source, "_leading_rows", namespace) scheduler_cls = _extract_class(source, "FusedCommMoEScheduler", {"_forward_chunk"}, namespace) - plan = object() + plan = SimpleNamespace(reserved_sms=4) layer = None if mode == "off" else SimpleNamespace(produces_replica_plan=mode != "standard") def route(routes, use_dp): @@ -2359,13 +2403,15 @@ def run_moe(ctx): assert out == "out" tail = [("quantize", {}), ("bucket", 6)] if mode == "per_layer_dynamic": - # The planner consumes the logical routes the calibrator records or replays. + # The planner consumes the logical routes the calibrator records or + # replays, and quantization leaves the plan's SMs to its copy kernel. expected = [ ("update_statistic", "logical"), ("statistic", "logical"), ("calibrator", "logical"), ("route", "replayed logical"), - *tail, + ("quantize", {"reserved_sms": 4}), + ("bucket", 6), ("run_moe", "physical", plan), ] elif mode == "standard": diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index 4bd4ac797be0..4e4b79cc69de 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -2147,32 +2147,142 @@ def test_megamoe_live_plane_spec_rejects_expand_before_allocation(): @pytest.mark.cpu_only -def test_megamoe_rejects_expand_before_provider_allocation(monkeypatch): - moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) - torch.nn.Module.__init__(moe) - moe._weights_created = False - moe.expert_size_per_partition = 4 - moe._rebalance_slots_active = 1 - moe._rebalance_home_experts = None - moe.layer_idx = 0 - moe.hidden_size = 160 - moe.intermediate_size_per_partition = 80 - # The base class derives this from the expand ratio; pin an inconsistent value. - monkeypatch.setattr(MegaMoECuteDsl, "expand_intermediate_size_per_partition", 176) - moe.quant_method = NVFP4MegaMoECuteDslMethod() - moe.ep_size = 2 - allocate_provider = MagicMock() - moe._alloc_symm_provider = allocate_provider +def test_megamoe_rejects_expand_before_arena_allocation(): allocate_arena = MagicMock() - moe.layer_load_balancer = SimpleNamespace(allocate_arena=allocate_arena) + module = SimpleNamespace( + hidden_size=160, + intermediate_size_per_partition=80, + expand_intermediate_size_per_partition=176, + layer_load_balancer=SimpleNamespace( + produces_replica_plan=True, allocate_arena=allocate_arena + ), + ) with pytest.raises(ValueError, match="expand_intermediate == 2"): - moe.create_weights() + NVFP4MegaMoECuteDslMethod().create_weights(module) - allocate_provider.assert_not_called() allocate_arena.assert_not_called() +@pytest.mark.cpu_only +@pytest.mark.parametrize("produces_replica_plan", [True, False, None]) +def test_megamoe_quant_method_allocates_and_binds_the_arena_through_the_layer_handle( + monkeypatch, produces_replica_plan +): + method = NVFP4MegaMoECuteDslMethod() + bundle = method._live_weight_bundle_layout(160, 80, 160) + resident, helpers = 2, 1 if produces_replica_plan else 0 + arena = object() + events = [] + + def create_parent_weights(self, module, *args): + events.append("parent") + for name in NVFP4MegaMoECuteDslMethod._STREAMED_SOURCE_PARAMS: + module.register_parameter( + name, torch.nn.Parameter(torch.empty(1, dtype=torch.uint8), requires_grad=False) + ) + for name in ("fc31_alpha", "fc2_alpha"): + shape, _, dtype_name, _, _ = method._live_weight_plane(bundle, name).storage_metadata( + resident + ) + module.register_parameter( + name, + torch.nn.Parameter( + torch.ones(shape, dtype=getattr(torch, dtype_name)), requires_grad=False + ), + ) + + def allocate_arena(**kwargs): + events.append(("allocate", kwargs)) + return arena + + monkeypatch.setattr(NVFP4FusedMoEMethod, "create_weights", create_parent_weights) + monkeypatch.setattr(NVFP4MegaMoECuteDslMethod, "setup_quant_scales", lambda self, module: None) + monkeypatch.setattr( + NVFP4MegaMoECuteDslMethod, + "_bind_live_weight_planes", + lambda self, module, bound: events.append(("bind", bound)), + ) + module = torch.nn.Module() + module.hidden_size = 160 + module.intermediate_size_per_partition = 80 + module.expand_intermediate_size_per_partition = 160 + module.bias = False + module.expert_size_per_partition = resident + module.compute_slot_count = resident + helpers + module.rebuild_tensor_metadata = {} + module.ep_process_group = "ep_group" + module.mapping = SimpleNamespace(local_rank=3) + module.layer_idx = 5 + module.layer_load_balancer = ( + None + if produces_replica_plan is None + else SimpleNamespace( + produces_replica_plan=produces_replica_plan, allocate_arena=allocate_arena + ) + ) + + method.create_weights(module) + + assert module.mega_fc1_weight.shape[0] == resident + helpers + if not produces_replica_plan: + assert events == ["parent"] + return + # Every rank allocates the arena before any Parameter exists and binds the + # seven planes once all of them are registered. + assert events == [ + ( + "allocate", + dict( + bundle=bundle, + ep_process_group="ep_group", + mapping=module.mapping, + device=3, + layer_idx=5, + ), + ), + "parent", + ("bind", arena), + ] + + +@pytest.mark.cpu_only +def test_megamoe_quant_method_builds_the_producer_after_checking_the_planes(monkeypatch): + events = [] + monkeypatch.setattr( + NVFP4FusedMoEMethod, "cache_derived_state", lambda self, module: events.append("parent") + ) + monkeypatch.setattr( + NVFP4MegaMoECuteDslMethod, + "assert_live_weight_aliases", + lambda self, module, arena: events.append(("aliases", arena)), + ) + arena = SimpleNamespace(assert_identity=lambda: events.append("identity")) + handle = SimpleNamespace( + produces_replica_plan=True, + live_weight_arena=arena, + bind_backend=lambda module: events.append(("bind_backend", module)), + ) + module = SimpleNamespace(layer_load_balancer=handle) + method = NVFP4MegaMoECuteDslMethod() + + method.cache_derived_state(module) + + assert events == ["parent", "identity", ("aliases", arena), ("bind_backend", module)] + + # Standard EPLB and EPLB-off layers only refresh the parent state. + events.clear() + method.cache_derived_state(SimpleNamespace(layer_load_balancer=None)) + method.cache_derived_state( + SimpleNamespace(layer_load_balancer=SimpleNamespace(produces_replica_plan=False)) + ) + assert events == ["parent", "parent"] + + handle.live_weight_arena = None + with pytest.raises(RuntimeError, match="no live-weight arena"): + method.cache_derived_state(module) + + @pytest.mark.cpu_only def test_megamoe_live_plane_identity_guard_detects_parameter_replacement(): class Tensor: @@ -2198,7 +2308,7 @@ def data_ptr(self): local_plane_views=(arena_view,), tekit_alias_views=(alias,), ) - module = SimpleNamespace(_rebalance_arena=arena, weight=Tensor(17)) + module = SimpleNamespace(weight=Tensor(17)) method = NVFP4MegaMoECuteDslMethod() method.live_weight_plane_spec = MagicMock(return_value=bundle) @@ -2209,32 +2319,33 @@ def data_ptr(self): method.assert_live_weight_aliases(module, arena) +def _per_layer_dynamic_handle(helper_slots_per_rank): + return SimpleNamespace(produces_replica_plan=True, helper_slots_per_rank=helper_slots_per_rank) + + @pytest.mark.cpu_only -def test_megamoe_compute_slots_do_not_widen_resident_expert_axis(): +def test_megamoe_compute_slots_come_from_the_layer_handle(): moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) torch.nn.Module.__init__(moe) moe.expert_size_per_partition = 8 - moe._rebalance_home_experts = 8 - moe._rebalance_slots_active = 3 + moe.layer_load_balancer = None + assert moe.compute_slot_count == 8 + # Standard EPLB produces no replica plan and adds no helper slot. + moe.layer_load_balancer = SimpleNamespace(produces_replica_plan=False) + assert moe.compute_slot_count == 8 + + moe.layer_load_balancer = _per_layer_dynamic_handle(3) assert moe.compute_slot_count == 11 assert moe.expert_size_per_partition == 8 - # A resident axis that already counts the helper slots - # (expert_size_per_partition == H + S) must raise instead of adding S again. - moe.expert_size_per_partition = 11 - with pytest.raises(RuntimeError, match="resident expert geometry changed"): - _ = moe.compute_slot_count - @pytest.mark.cpu_only def test_megamoe_replica_plan_supplies_launch_facts_and_route_wait(): moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) torch.nn.Module.__init__(moe) moe.expert_size_per_partition = 8 - moe._rebalance_home_experts = 8 - moe._rebalance_slots_active = 3 - moe._quantize_reserved_sms = 8 + moe.layer_load_balancer = _per_layer_dynamic_handle(3) wait_for_routes = MagicMock() ready_flags = torch.empty(8, dtype=torch.int64) plan = MoEReplicaPlan( @@ -2242,14 +2353,14 @@ def test_megamoe_replica_plan_supplies_launch_facts_and_route_wait(): compute_slots_per_rank=11, ready_flags=ready_flags, ready_generation=7, - reserved_sms=8, + reserved_sms=4, _wait_for_routes=wait_for_routes, ) compute_slots, helper_slots, actual_flags, generation, reserved_sms = ( moe._replica_plan_launch_args(plan) ) - assert (compute_slots, helper_slots, generation, reserved_sms) == (11, 3, 7, 8) + assert (compute_slots, helper_slots, generation, reserved_sms) == (11, 3, 7, 4) assert actual_flags is ready_flags moe._wait_rebalance_routes(plan) wait_for_routes.assert_called_once_with() @@ -2344,14 +2455,52 @@ def test_megamoe_run_moe_hands_routes_to_staging_unconverted(): assert torch.equal(captured["topk_weights"], scales) +@pytest.mark.cpu_only +@pytest.mark.parametrize("has_sm_budget", [True, False]) +def test_megamoe_quantize_reserves_sms_only_for_a_replica_plan(monkeypatch, request, has_sm_budget): + from tensorrt_llm._torch.moe.fused_moe.mega_moe import mega_moe_cute_dsl + + calls = [] + + def outputs(x): + rows = x.shape[0] + return ( + torch.zeros(rows, x.shape[1] // 2, dtype=torch.uint8), + torch.zeros(rows * x.shape[1] // 16, dtype=torch.uint8), + ) + + def fp4_quantize(x, scale, sf_vec_size, sf_use_ue8m0, is_sf_swizzled): + calls.append(("default", sf_vec_size, sf_use_ue8m0, is_sf_swizzled)) + return outputs(x) + + def sm_budget(x, scale, sf_vec_size, sf_use_ue8m0, is_sf_swizzled, *, reserved_sms): + calls.append(("sm_budget", sf_vec_size, sf_use_ue8m0, is_sf_swizzled, reserved_sms)) + return outputs(x) + + if has_sm_budget: + fp4_quantize.sm_budget = sm_budget + monkeypatch.setattr(torch.ops.trtllm, "fp4_quantize", fp4_quantize, raising=False) + mega_moe_cute_dsl._fp4_quantize_has_sm_budget.cache_clear() + request.addfinalizer(mega_moe_cute_dsl._fp4_quantize_has_sm_budget.cache_clear) + moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) + torch.nn.Module.__init__(moe) + moe.fc31_input_scale = torch.ones(1) + x = torch.zeros(3, 64, dtype=torch.bfloat16) + + moe.quantize_input(x) + moe.quantize_input(x, reserved_sms=4) + + default = ("default", 16, False, False) + budget = ("sm_budget", 16, False, False, 4) if has_sm_budget else default + assert calls == [default, budget] + + @pytest.mark.cpu_only def test_megamoe_replica_plan_rejects_unallocated_geometry(): moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) torch.nn.Module.__init__(moe) moe.expert_size_per_partition = 8 - moe._rebalance_home_experts = 8 - moe._rebalance_slots_active = 3 - moe._quantize_reserved_sms = 8 + moe.layer_load_balancer = _per_layer_dynamic_handle(3) plan = MoEReplicaPlan( resident_experts_per_rank=8, compute_slots_per_rank=12, @@ -2372,24 +2521,14 @@ def test_megamoe_cutedsl_post_load_weights_uses_staged_hooks(): quant_method = SimpleNamespace( transform_weights=MagicMock(), cache_derived_state=MagicMock(), - assert_live_weight_aliases=MagicMock(), ) moe.quant_method = quant_method - moe._rebalance_slots_active = 1 - arena = MagicMock() - # ``unittest.mock`` blocks dynamically-created ``assert_*`` attributes. - arena.assert_identity = MagicMock() - moe._rebalance_arena = arena - moe.layer_load_balancer = SimpleNamespace(bind_backend=MagicMock()) moe.post_load_weights() moe.transform_weights() quant_method.transform_weights.assert_called_once_with(moe) quant_method.cache_derived_state.assert_called_once_with(moe) - moe._rebalance_arena.assert_identity.assert_called_once_with() - quant_method.assert_live_weight_aliases.assert_called_once_with(moe, moe._rebalance_arena) - moe.layer_load_balancer.bind_backend.assert_called_once_with(moe) assert moe._weights_transformed is True @@ -2442,6 +2581,7 @@ def test_megamoe_cutedsl_cache_derived_state_reuses_rebalance_binding( PerLayerDynamicMoeLoadBalancer, planner, ) + from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb import arena as arena_module config = SimpleNamespace( num_slots=24, @@ -2458,21 +2598,39 @@ def test_megamoe_cutedsl_cache_derived_state_reuses_rebalance_binding( top_k=2, slot_count_per_rank=6, ) - arena = MagicMock(home_experts=4, helper_slots=2) + bundle = NVFP4MegaMoECuteDslMethod._live_weight_bundle_layout(160, 80, 160) + arena = MagicMock(home_experts=4, helper_slots=2, bundle=bundle) arena.assert_identity = MagicMock() + allocate = MagicMock(return_value=arena) + monkeypatch.setattr(arena_module, "allocate_rebalance_arena", allocate) + monkeypatch.setattr(manager, "get_ep_comm", lambda process_group: "ep_comm") + monkeypatch.setattr(manager, "create_arena_provider", lambda: "provider") moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) torch.nn.Module.__init__(moe) - moe.quant_method = SimpleNamespace( - cache_derived_state=MagicMock(), - assert_live_weight_aliases=MagicMock(), - ) - moe._rebalance_slots_active = 2 - moe._rebalance_arena = arena moe._ep_pg = object() moe._maxt_buckets = [64, 128] moe.mapping = SimpleNamespace(local_rank=0) moe.layer_idx = 5 moe.layer_load_balancer = layer + moe.initial_global_assignments = list(range(16)) + # The quant method's create_weights allocates through the handle; a second + # call with the same plane spec returns the same arena. + for _ in range(2): + assert ( + layer.allocate_arena( + bundle=bundle, + ep_process_group=moe.ep_process_group, + mapping=moe.mapping, + device=0, + layer_idx=5, + ) + is arena + ) + allocate.assert_called_once() + moe.quant_method = NVFP4MegaMoECuteDslMethod() + monkeypatch.setattr(NVFP4MegaMoECuteDslMethod, "setup_quant_scales", lambda self, module: None) + assert_aliases = MagicMock() + monkeypatch.setattr(NVFP4MegaMoECuteDslMethod, "assert_live_weight_aliases", assert_aliases) builder = MagicMock(return_value=SimpleNamespace(home_experts=4, helper_slots=2)) monkeypatch.setattr( planner, @@ -2483,9 +2641,9 @@ def test_megamoe_cutedsl_cache_derived_state_reuses_rebalance_binding( moe.cache_derived_state() moe.cache_derived_state() - assert moe.quant_method.cache_derived_state.call_count == 2 assert arena.assert_identity.call_count == 2 - assert moe.quant_method.assert_live_weight_aliases.call_count == 2 + assert assert_aliases.call_count == 2 + assert layer.live_weight_arena is arena builder.assert_called_once_with( arena=arena, manager=manager, diff --git a/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py b/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py index 5c0f28b30284..1e9836aa00be 100644 --- a/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py +++ b/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py @@ -46,7 +46,8 @@ class _Arena: home_experts = 4 helper_slots = 2 - def __init__(self): + def __init__(self, bundle="bundle"): + self.bundle = bundle self.flags = torch.empty(4, dtype=torch.uint64) self.generation = 0 self.provider = SimpleNamespace(pool=None) @@ -62,8 +63,7 @@ def assert_identity(self): class _Backend: - def __init__(self, arena): - self.live_weight_arena = arena + def __init__(self): self.tactic_autotune = False self.mapping = SimpleNamespace(local_rank=0) self.ep_process_group = object() @@ -115,6 +115,22 @@ def _manager(): ) +def _allocate_arena(layer, monkeypatch, arena): + """Allocate ``arena`` through the layer handle, as the quant method does.""" + from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb import arena as arena_module + + monkeypatch.setattr(layer.manager, "get_ep_comm", lambda process_group: object()) + monkeypatch.setattr(layer.manager, "create_arena_provider", lambda: object()) + monkeypatch.setattr(arena_module, "allocate_rebalance_arena", MagicMock(return_value=arena)) + return layer.allocate_arena( + bundle=arena.bundle, + ep_process_group=object(), + mapping=object(), + device=0, + layer_idx=layer.layer_idx, + ) + + def test_layer_capacity_uses_configured_compute_slots(): manager = _manager() with pytest.raises(ValueError, match="configured local slot count"): @@ -247,7 +263,8 @@ def test_bind_backend_is_idempotent_before_scheduler_construction(monkeypatch): manager = _manager() layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) arena = _Arena() - backend = _Backend(arena) + assert _allocate_arena(layer, monkeypatch, arena) is arena + backend = _Backend() group = _Group() builder = MagicMock(return_value=group) monkeypatch.setattr( @@ -276,30 +293,34 @@ def test_bind_backend_is_idempotent_before_scheduler_construction(monkeypatch): assert layer._scheduler_group is group -def test_bind_backend_rejects_different_backend_or_arena_before_construction( +def test_bind_backend_requires_the_layer_arena_and_rejects_another_backend( monkeypatch, ): from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb import planner manager = _manager() layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) - arena = _Arena() - backend = _Backend(arena) + backend = _Backend() builder = MagicMock(return_value=_Group()) monkeypatch.setattr( planner, "build_rebalance_slot_scheduler_group", builder, ) + with pytest.raises(RuntimeError, match="no live-weight arena"): + layer.bind_backend(backend) + + arena = _Arena() + _allocate_arena(layer, monkeypatch, arena) + # The runtime binds only the arena this handle allocated. + with pytest.raises(RuntimeError, match="did not allocate"): + layer.bind_runtime(arena=_Arena(), scheduler_group=_Group()) layer.bind_backend(backend) with pytest.raises(RuntimeError, match="different backend"): - layer.bind_backend(_Backend(arena)) - backend.live_weight_arena = _Arena() - with pytest.raises(RuntimeError, match="changed its rebalance arena"): - layer.bind_backend(backend) - + layer.bind_backend(_Backend()) builder.assert_called_once() + assert layer.live_weight_arena is arena def test_iteration_context_reuses_common_lifecycle(): @@ -380,7 +401,8 @@ def test_layer_allocates_arena_from_model_owned_resources(monkeypatch): manager = _manager() layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) order = [] - process_group, comm, provider, arena = object(), object(), object(), object() + process_group, comm, provider, bundle, mapping = (object() for _ in range(5)) + arena = SimpleNamespace(bundle=bundle) def get_ep_comm(group): assert group is process_group @@ -395,17 +417,12 @@ def create_arena_provider(): monkeypatch.setattr(manager, "create_arena_provider", create_arena_provider) allocate = MagicMock(return_value=arena) monkeypatch.setattr(arena_module, "allocate_rebalance_arena", allocate) - bundle, mapping = object(), object() + kwargs = dict(ep_process_group=process_group, mapping=mapping, device=1, layer_idx=7) - allocated = layer.allocate_arena( - bundle=bundle, - ep_process_group=process_group, - mapping=mapping, - device=1, - layer_idx=7, - ) + allocated = layer.allocate_arena(bundle=bundle, **kwargs) assert allocated is arena + assert layer.live_weight_arena is arena assert order == ["ep_comm", "provider"] allocate.assert_called_once_with( home_experts=4, @@ -418,6 +435,13 @@ def create_arena_provider(): layer_idx=7, ) + # Allocation happens once per layer; only the same plane spec gets the arena back. + assert layer.allocate_arena(bundle=bundle, **kwargs) is arena + with pytest.raises(RuntimeError, match="different live-weight plane spec"): + layer.allocate_arena(bundle=object(), **kwargs) + allocate.assert_called_once() + assert order == ["ep_comm", "provider"] + @pytest.mark.parametrize( ("helper_slot_mode", "per_layer_sets"), From 3d4f964aa71a4cd9b33519c45c25caf2eacfbcad Mon Sep 17 00:00:00 2001 From: jintaop Date: Fri, 23 Oct 2026 18:56:56 -0700 Subject: [PATCH 37/59] [None][test] Follow main's engine warmup and release API in the dynamic EPLB contract tests Main now applies the engine's warmup effects in warmup_scope() through _set_warmup_effects(), and PyExecutor releases each engine's graphs through release_executor_resources(). The warmup contract test drives the scope and checks both transitions; the shutdown-retry test's fake engine implements the release method the executor now calls. Signed-off-by: jintaop --- .../_torch/moe/test_dynamic_eplb_contracts.py | 24 ++++++++++++------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index b16ed958c30b..e85edc2b43c5 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -13,7 +13,7 @@ import re import sys import threading -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext from pathlib import Path from types import ModuleType, SimpleNamespace from typing import Any, List, Optional, Tuple @@ -591,8 +591,11 @@ def test_model_engine_warmup_uses_load_balancer_public_api(): engine_cls = _extract_class( _TORCH_ROOT / "pyexecutor/model_engine.py", "PyTorchModelEngine", - {"is_warmup", "moe_load_balancer_iter_info"}, - {"set_moe_comm_warmup": lambda value: events.append(("comm", value))}, + {"_set_warmup_effects", "warmup_scope"}, + { + "contextmanager": contextmanager, + "set_moe_comm_warmup": lambda value: events.append(("comm", value)), + }, ) manager = SimpleNamespace( enable_statistic=True, @@ -601,12 +604,17 @@ def test_model_engine_warmup_uses_load_balancer_public_api(): ) engine = engine_cls() engine.moe_load_balancer = manager + engine._warmup_depth = 0 - engine.is_warmup = True + with engine.warmup_scope(): + assert events == [ + ("comm", True), + ("iter", {"enable_statistic": False, "enable_update_weights": False}), + ] - assert events == [ - ("comm", True), - ("iter", {"enable_statistic": False, "enable_update_weights": False}), + assert events[2:] == [ + ("comm", False), + ("iter", {"enable_statistic": True, "enable_update_weights": True}), ] @@ -842,7 +850,7 @@ def async_worker_stop(self): raise RuntimeError("sampler stop failed") class Engine: - def _release_cuda_graphs(self, **_): + def release_executor_resources(self, *, release_model_owned_graphs): events.append("release_graphs") class DwdpManager: From 5c283d2435c979837246066260dbe55acb5bcb5c Mon Sep 17 00:00:00 2001 From: jintaop Date: Sun, 25 Oct 2026 06:46:31 -0700 Subject: [PATCH 38/59] [None][perf] Submit per-layer dynamic EPLB with one native call per layer Each MoE layer's HALO-Q + TMA launch took two cuda-python event calls, a Python stream switch and a 25-argument operator call, and the consumer release two more cuda-python calls. Passing the TorchScript custom-class context to a torch.ops call alone costs about 8 us of dispatch. On VR200 with main's SM107 CuteDSL attention GEMMs, rank 0's executor host paces the ON arm, so this per-layer host time showed up end to end. Bind every argument that the fused launches share, the route epoch and the copy stream into the native TMA context once. The per-layer ops take an int handle of that context: moe_rebalance_tma_submit_fused passes only the routes, the route target and the spin bound, and orders the copy stream after the current stream itself; moe_rebalance_tma_release does the same for the consumer of the live bank. The endpoint that owns the context keeps it alive and clears the handle before destroying it. A source contract pins the native call order. Signed-off-by: jintaop --- .../moe/loadBalance/moeRebalanceHaloQOp.h | 2 +- .../moe/loadBalance/moeRebalanceTmaOp.cpp | 208 ++++++++++++-- .../per_layer_dynamic_eplb/planner.py | 37 +-- .../cuda_scheduler/runtime.py | 98 ++----- .../megamoe_scheduler/sami/hierarchical.py | 76 +++-- .../_torch/moe/test_dynamic_eplb_contracts.py | 271 ++++++++++-------- tests/unittest/_torch/moe/test_moe_backend.py | 19 +- 7 files changed, 422 insertions(+), 289 deletions(-) diff --git a/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.h b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.h index d3bef11930b2..ef9e075da699 100644 --- a/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.h +++ b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.h @@ -27,7 +27,7 @@ TRTLLM_NAMESPACE_BEGIN namespace torch_ext { -// Validates the HALO-Q scheduler arguments of moe_rebalance_tma_submit_fused and returns its launch parameters. +// Validates the HALO-Q scheduler arguments that moe_rebalance_tma_bind_fused binds and returns its launch parameters. kernels::MoeRebalanceHaloQParams makeMoeRebalanceHaloQParams(torch::Tensor const& routes, torch::Tensor& outSlots, torch::Tensor& outIds, torch::Tensor& outLevels, torch::Tensor& outOwners, torch::Tensor const& peerBases, torch::Tensor& status, torch::Tensor& partial, torch::Tensor& routeAux, torch::Tensor& gridSync, diff --git a/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp index 39f34ca6bc74..4a43b222ec98 100644 --- a/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp +++ b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp @@ -43,6 +43,11 @@ void checkTmaResult(int error, char const* operation) TORCH_CHECK(error == 0, operation, " failed: ", megamoe_tma_copy_error_string(error), " (", error, ")"); } +void checkCudaResult(cudaError_t error, char const* operation) +{ + TORCH_CHECK(error == cudaSuccess, operation, " failed: ", cudaGetErrorString(error)); +} + int checkedInt(int64_t value, char const* name) { TORCH_CHECK(value >= std::numeric_limits::min() && value <= std::numeric_limits::max(), name, @@ -115,6 +120,7 @@ class MoeRebalanceTmaState : public torch::CustomClassHolder // available when the caller needs a synchronous error result. static_cast(megamoe_tma_copy_destroy(&mState)); } + releaseFusedEvents(); } void configureGpuPlan(int64_t world, int64_t rank, int64_t helperCount, int64_t planes, int64_t globalExperts, @@ -229,12 +235,17 @@ class MoeRebalanceTmaState : public torch::CustomClassHolder mBound = true; } - // Launches the bound HALO-Q scheduler and this copy as one kernel; see megamoe_tma_copy_submit_fused. - void submitFused(int64_t flagMc, kernels::MoeRebalanceHaloQParams const& scheduler, torch::Tensor const& routeEpoch, - int64_t routeTarget) + // Binds every argument that the fused launches share: the HALO-Q scheduler workspace and outputs, the route + // epoch, and the static copy stream. A launch then passes only its routes. + void bindFused(int64_t flagMc, torch::Tensor routeEpoch, torch::Tensor outSlots, torch::Tensor outIds, + torch::Tensor outLevels, torch::Tensor outOwners, torch::Tensor peerBases, torch::Tensor status, + torch::Tensor partial, torch::Tensor routeAux, torch::Tensor gridSync, torch::Tensor planWorkspace, + torch::Tensor routePrefix, int64_t ep, int64_t experts, int64_t helpers, int64_t localRank, int64_t topk, + int64_t routeCapacity, int64_t ctas, int64_t threads, int64_t copyStream) { std::lock_guard lock(mMutex); requireBound(); + TORCH_CHECK(!mFusedBound, "moe rebalance fused launch is already bound"); TORCH_CHECK( routeEpoch.is_cuda() && routeEpoch.get_device() == mDevice, "route_epoch must be on CUDA device ", mDevice); TORCH_CHECK( @@ -242,17 +253,84 @@ class MoeRebalanceTmaState : public torch::CustomClassHolder "route_epoch must be a contiguous one-element torch.int64 tensor"); uint64_t const flag = checkedUnsigned(flagMc, "flag_mc"); TORCH_CHECK(flag != 0, "flag_mc must be nonzero"); - kernels::checkMoeRebalanceHaloQParams(scheduler); + TORCH_CHECK(topk > 0 && routeCapacity % topk == 0, "route_capacity must be a positive multiple of topk"); + TORCH_CHECK(copyStream != 0, "copy_stream must be a non-default CUDA stream"); + TORCH_CHECK( + outSlots.is_cuda() && outSlots.get_device() == mDevice, "out_slots must be on CUDA device ", mDevice); + // Validate the shared arguments once, with no routes; a launch only adds its routes. + auto params = makeMoeRebalanceHaloQParams(torch::empty({0, topk}, outSlots.options()), outSlots, outIds, + outLevels, outOwners, peerBases, status, partial, routeAux, gridSync, planWorkspace, routePrefix, ep, + experts, helpers, localRank, routeCapacity, ctas, threads, 0, 0); + kernels::checkMoeRebalanceHaloQParams(params); + + c10::cuda::CUDAGuard const deviceGuard(c10::Device(c10::DeviceType::CUDA, mDevice)); + cudaEvent_t inputReady{nullptr}; + checkCudaResult( + cudaEventCreateWithFlags(&inputReady, cudaEventDisableTiming), "cudaEventCreate(rebalance input ready)"); + cudaEvent_t consumerDone{nullptr}; + cudaError_t const error = cudaEventCreateWithFlags(&consumerDone, cudaEventDisableTiming); + if (error != cudaSuccess) + { + static_cast(cudaEventDestroy(inputReady)); + checkCudaResult(error, "cudaEventCreate(rebalance consumer done)"); + } + + mRouteEpochPtr = reinterpret_cast(routeEpoch.data_ptr()); + // The fused launch reads and writes these by address; retaining them prevents allocator reuse. + mFusedTensors = {std::move(routeEpoch), std::move(outSlots), std::move(outIds), std::move(outLevels), + std::move(outOwners), std::move(peerBases), std::move(status), std::move(partial), std::move(routeAux), + std::move(gridSync), std::move(planWorkspace), std::move(routePrefix)}; + mFusedParams = params; + mFlagMc = flag; + mTopk = topk; + mCopyStream = reinterpret_cast(copyStream); + mInputReady = inputReady; + mConsumerDone = consumerDone; + mFusedBound = true; + } + + // Launches the bound HALO-Q scheduler and this copy as one kernel on the copy stream, ordered after the work + // already queued on the current stream; see megamoe_tma_copy_submit_fused. + void submitFused(torch::Tensor const& routes, int64_t routeTarget, int64_t spinCycles) + { + std::lock_guard lock(mMutex); + requireFusedBound(); + TORCH_CHECK_VALUE(routes.is_cuda() && routes.get_device() == mDevice + && routes.scalar_type() == at::ScalarType::Int && routes.is_contiguous() && routes.dim() == 2 + && routes.size(1) == mTopk && routes.numel() <= mFusedParams.routeCapacity, + "routes must be a contiguous CUDA int32 [tokens, ", mTopk, "] tensor on device ", mDevice, " with at most ", + mFusedParams.routeCapacity / mTopk, " tokens"); + TORCH_CHECK_VALUE(spinCycles > 0, "spin_cycles must be positive"); + auto scheduler = mFusedParams; + scheduler.routes = routes.data_ptr(); + scheduler.validRouteCount = static_cast(routes.numel()); + scheduler.spinCycles = static_cast(spinCycles); + c10::cuda::CUDAGuard const deviceGuard(c10::Device(c10::DeviceType::CUDA, mDevice)); - auto const stream = at::cuda::getCurrentCUDAStream(mDevice).stream(); - checkTmaResult(megamoe_tma_copy_submit_fused(mState, scheduler, - reinterpret_cast(routeEpoch.data_ptr()), - static_cast(routeTarget), flag, mGeneration, reinterpret_cast(stream)), + cudaStream_t const execution = currentExecutionStream(); + checkCudaResult(cudaEventRecord(mInputReady, execution), "cudaEventRecord(rebalance input ready)"); + checkCudaResult(cudaStreamWaitEvent(mCopyStream, mInputReady, 0), "cudaStreamWaitEvent(rebalance input ready)"); + checkTmaResult( + megamoe_tma_copy_submit_fused(mState, scheduler, mRouteEpochPtr, static_cast(routeTarget), + mFlagMc, mGeneration, reinterpret_cast(mCopyStream)), "moe rebalance fused HALO-Q + TMA submit"); ++mGeneration; TORCH_CHECK(mGeneration != 0, "moe rebalance TMA generation overflow"); } + // Orders the copy stream after the work already queued on the current stream, such as the consumer of the + // live bank that the next fused launch overwrites. + void releaseAfterCurrentStream() + { + std::lock_guard lock(mMutex); + requireFusedBound(); + c10::cuda::CUDAGuard const deviceGuard(c10::Device(c10::DeviceType::CUDA, mDevice)); + cudaStream_t const execution = currentExecutionStream(); + checkCudaResult(cudaEventRecord(mConsumerDone, execution), "cudaEventRecord(rebalance consumer done)"); + checkCudaResult( + cudaStreamWaitEvent(mCopyStream, mConsumerDone, 0), "cudaStreamWaitEvent(rebalance consumer done)"); + } + int64_t currentGeneration() const { std::lock_guard lock(mMutex); @@ -312,6 +390,11 @@ class MoeRebalanceTmaState : public torch::CustomClassHolder mGridSync = torch::Tensor(); mConfigured = false; mBound = false; + releaseFusedEvents(); + mFusedTensors.clear(); + mRouteEpochPtr = nullptr; + mCopyStream = nullptr; + mFusedBound = false; } checkTmaResult(error, "moe rebalance TMA destroy"); } @@ -328,6 +411,35 @@ class MoeRebalanceTmaState : public torch::CustomClassHolder TORCH_CHECK(mBound, "moe rebalance TMA state must be bound before submit"); } + void requireFusedBound() const + { + requireBound(); + TORCH_CHECK(mFusedBound, "moe rebalance fused launch must be bound before submit"); + } + + // The stream whose queued work a launch or release follows; it must differ from the copy stream. + cudaStream_t currentExecutionStream() const + { + cudaStream_t const stream = at::cuda::getCurrentCUDAStream(mDevice).stream(); + TORCH_CHECK(stream != mCopyStream, "the current stream must differ from the bound copy stream"); + return stream; + } + + void releaseFusedEvents() noexcept + { + // Pending records complete normally after an event is destroyed. + if (mInputReady != nullptr) + { + static_cast(cudaEventDestroy(mInputReady)); + mInputReady = nullptr; + } + if (mConsumerDone != nullptr) + { + static_cast(cudaEventDestroy(mConsumerDone)); + mConsumerDone = nullptr; + } + } + static int64_t checkedDiagnostic(uint64_t value) { TORCH_CHECK(value <= static_cast(std::numeric_limits::max()), @@ -348,6 +460,16 @@ class MoeRebalanceTmaState : public torch::CustomClassHolder torch::Tensor mWorkspace; torch::Tensor mStatus; torch::Tensor mGridSync; + // Arguments bound by bindFused() and shared by every fused launch. + kernels::MoeRebalanceHaloQParams mFusedParams{}; + std::vector mFusedTensors; + uint64_t* mRouteEpochPtr{nullptr}; + uint64_t mFlagMc{0}; + int64_t mTopk{0}; + cudaStream_t mCopyStream{nullptr}; + cudaEvent_t mInputReady{nullptr}; + cudaEvent_t mConsumerDone{nullptr}; + bool mFusedBound{false}; }; using MoeRebalanceTmaStatePtr = c10::intrusive_ptr; @@ -379,18 +501,47 @@ void moeRebalanceTmaBindGpuDirect(MoeRebalanceTmaStatePtr const& state, torch::T std::move(gridSync), capacity, helperResidency); } -void moeRebalanceTmaSubmitFused(MoeRebalanceTmaStatePtr const& state, int64_t flagMc, torch::Tensor& routeEpoch, - int64_t routeTarget, torch::Tensor const& routes, torch::Tensor& outSlots, torch::Tensor& outIds, - torch::Tensor& outLevels, torch::Tensor& outOwners, torch::Tensor const& peerBases, torch::Tensor& status, - torch::Tensor& partial, torch::Tensor& routeAux, torch::Tensor& gridSync, torch::Tensor& planWorkspace, - torch::Tensor& routePrefix, int64_t ep, int64_t experts, int64_t helpers, int64_t localRank, int64_t routeCapacity, - int64_t ctas, int64_t threads, int64_t spinCycles, int64_t validRouteCount) +void moeRebalanceTmaBindFused(MoeRebalanceTmaStatePtr const& state, int64_t flagMc, torch::Tensor routeEpoch, + torch::Tensor outSlots, torch::Tensor outIds, torch::Tensor outLevels, torch::Tensor outOwners, + torch::Tensor peerBases, torch::Tensor status, torch::Tensor partial, torch::Tensor routeAux, + torch::Tensor gridSync, torch::Tensor planWorkspace, torch::Tensor routePrefix, int64_t ep, int64_t experts, + int64_t helpers, int64_t localRank, int64_t topk, int64_t routeCapacity, int64_t ctas, int64_t threads, + int64_t copyStream) { TORCH_CHECK(state, "moe rebalance TMA state is null"); - auto const scheduler = makeMoeRebalanceHaloQParams(routes, outSlots, outIds, outLevels, outOwners, peerBases, - status, partial, routeAux, gridSync, planWorkspace, routePrefix, ep, experts, helpers, localRank, routeCapacity, - ctas, threads, spinCycles, validRouteCount); - state->submitFused(flagMc, scheduler, routeEpoch, routeTarget); + state->bindFused(flagMc, std::move(routeEpoch), std::move(outSlots), std::move(outIds), std::move(outLevels), + std::move(outOwners), std::move(peerBases), std::move(status), std::move(partial), std::move(routeAux), + std::move(gridSync), std::move(planWorkspace), std::move(routePrefix), ep, experts, helpers, localRank, topk, + routeCapacity, ctas, threads, copyStream); +} + +// The per-layer launch and release ops take this handle instead of the custom class, whose conversion costs several +// microseconds of host time per operator call. The caller keeps the state alive while it uses the handle. +int64_t moeRebalanceTmaHandle(MoeRebalanceTmaStatePtr const& state) +{ + TORCH_CHECK(state, "moe rebalance TMA state is null"); + return reinterpret_cast(state.get()); +} + +namespace +{ + +MoeRebalanceTmaState& stateFromHandle(int64_t handle) +{ + TORCH_CHECK(handle != 0, "moe rebalance TMA handle is null"); + return *reinterpret_cast(handle); +} + +} // namespace + +void moeRebalanceTmaSubmitFused(int64_t handle, torch::Tensor const& routes, int64_t routeTarget, int64_t spinCycles) +{ + stateFromHandle(handle).submitFused(routes, routeTarget, spinCycles); +} + +void moeRebalanceTmaRelease(int64_t handle) +{ + stateFromHandle(handle).releaseAfterCurrentStream(); } int64_t moeRebalanceTmaCurrentGeneration(MoeRebalanceTmaStatePtr const& state) @@ -439,13 +590,17 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state, Tensor ids, Tensor levels, Tensor owners, " "Tensor workspace, Tensor status, Tensor grid_sync, int capacity, bool helper_residency) -> ()"); m.def( - "moe_rebalance_tma_submit_fused(" - "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state, int flag_mc, Tensor(a!) route_epoch, " - "int route_target, Tensor routes, Tensor(b!) out_slots, Tensor(c!) out_ids, Tensor(d!) out_levels, " - "Tensor(e!) out_owners, Tensor peer_bases, Tensor(f!) status, Tensor(g!) partial, Tensor(h!) route_aux, " - "Tensor(i!) grid_sync, Tensor(j!) plan_workspace, Tensor(k!) route_prefix, int ep, int experts, " - "int helpers, int local_rank, int route_capacity, int ctas, int threads, int spin_cycles, " - "int valid_route_count) -> ()"); + "moe_rebalance_tma_bind_fused(" + "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state, int flag_mc, Tensor route_epoch, " + "Tensor out_slots, Tensor out_ids, Tensor out_levels, Tensor out_owners, Tensor peer_bases, Tensor status, " + "Tensor partial, Tensor route_aux, Tensor grid_sync, Tensor plan_workspace, Tensor route_prefix, int ep, " + "int experts, int helpers, int local_rank, int topk, int route_capacity, int ctas, int threads, " + "int copy_stream) -> ()"); + m.def( + "moe_rebalance_tma_handle(" + "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state) -> int"); + m.def("moe_rebalance_tma_submit_fused(int handle, Tensor routes, int route_target, int spin_cycles) -> ()"); + m.def("moe_rebalance_tma_release(int handle) -> ()"); m.def( "moe_rebalance_tma_current_gen(" "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state) -> int"); @@ -465,7 +620,10 @@ TORCH_LIBRARY_IMPL(trtllm, CompositeExplicitAutograd, m) m.impl("moe_rebalance_tma_create", &tensorrt_llm::torch_ext::moeRebalanceTmaCreate); m.impl("moe_rebalance_tma_configure_gpu_plan", &tensorrt_llm::torch_ext::moeRebalanceTmaConfigureGpuPlan); m.impl("moe_rebalance_tma_bind_gpu_direct", &tensorrt_llm::torch_ext::moeRebalanceTmaBindGpuDirect); + m.impl("moe_rebalance_tma_bind_fused", &tensorrt_llm::torch_ext::moeRebalanceTmaBindFused); + m.impl("moe_rebalance_tma_handle", &tensorrt_llm::torch_ext::moeRebalanceTmaHandle); m.impl("moe_rebalance_tma_submit_fused", &tensorrt_llm::torch_ext::moeRebalanceTmaSubmitFused); + m.impl("moe_rebalance_tma_release", &tensorrt_llm::torch_ext::moeRebalanceTmaRelease); m.impl("moe_rebalance_tma_current_gen", &tensorrt_llm::torch_ext::moeRebalanceTmaCurrentGeneration); m.impl("moe_rebalance_tma_result", &tensorrt_llm::torch_ext::moeRebalanceTmaResult); m.impl("moe_rebalance_tma_config", &tensorrt_llm::torch_ext::moeRebalanceTmaConfig); diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py index 9e0205adc7ba..72f0c8f186c1 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py @@ -205,15 +205,6 @@ def __init__( self.scheduler_stream = self.copy_stream self._stream_handle = int(self.copy_stream.cuda_stream) self._driver = driver - self._input_ready = torch.cuda.Event() - self._consumer_done = torch.cuda.Event() - # Materialize each lazily created event once. The steady path uses its - # cached driver handle, without per-plan stream guards or allocation. - with torch.cuda.device(self.device): - for event in (self._input_ready, self._consumer_done): - event.record(self.copy_stream) - self._input_ready_handle = int(self._input_ready.cuda_event) - self._consumer_done_handle = int(self._consumer_done.cuda_event) error, cu_device = driver.cuDeviceGet(self.device) supported = 0 if error == driver.CUresult.CUDA_SUCCESS: @@ -268,7 +259,7 @@ def __init__( tma_route="plan", helper_residency=bool(helper_residency), ) - self.broadcaster.bind_scheduler(self.scheduler, self.scheduler_stream) + self.broadcaster.bind_scheduler(self.scheduler, self.scheduler_stream, self._route_epoch) # A prior consumer event precedes the next HALO-Q EP rendezvous. self.lease = _LiveBankLeaseProvider(self) @@ -316,16 +307,6 @@ def _check_owner(self) -> None: raise RuntimeError("MAIN must be distinct from the higher-priority COPY stream") self._execution_stream_handle = handle - def _record_event(self, event_handle: int, stream_handle: int) -> None: - (error,) = self._driver.cuEventRecord(event_handle, stream_handle) - if error != self._driver.CUresult.CUDA_SUCCESS: - raise RuntimeError(f"cuEventRecord(rebalance) failed: {error!r}") - - def _wait_event(self, stream_handle: int, event_handle: int) -> None: - (error,) = self._driver.cuStreamWaitEvent(stream_handle, event_handle, 0) - if error != self._driver.CUresult.CUDA_SUCCESS: - raise RuntimeError(f"cuStreamWaitEvent(rebalance) failed: {error!r}") - def _wait_route_epoch(self, stream_handle: int) -> None: (error,) = self._driver.cuStreamWaitValue64( stream_handle, @@ -367,10 +348,6 @@ def plan_schedule(self, logical_expert_ids: torch.Tensor) -> Tuple[Any, int]: if num_tokens > self.max_tokens_per_rank: raise RuntimeError(f"Route chunk {num_tokens} exceeds {self.max_tokens_per_rank}") - # Both waits capture their event's current record before a later plan - # can re-record it. No host synchronization or temporary event is needed. - self._record_event(self._input_ready_handle, self._execution_stream_handle) - self._wait_event(self._stream_handle, self._input_ready_handle) outputs = self.scheduler.outputs self._scheduled_generation = self._generation + 1 if self._pending_release is not None: @@ -378,10 +355,11 @@ def plan_schedule(self, logical_expert_ids: torch.Tensor) -> Tuple[Any, int]: self._pending_release = None route_target = self._route_target + self.auxiliary_sms self._routes_in_flight = logical_expert_ids + # The launch follows the MAIN work queued so far, which wrote the + # routes; no host synchronization is needed. ticket = self.broadcaster.submit_fused( outputs, logical_expert_ids, - self._route_epoch, route_target, spin_cycles=self._spin_bound.cycles(), ) @@ -437,7 +415,11 @@ def discard_plan(self, part: Tuple[Any, int]) -> None: self.finish() def finish(self, completion_event: torch.cuda.Event | None = None) -> None: - """Queue the previous consumer wait before any next-generation producer.""" + """Queue the previous consumer wait before any next-generation producer. + + The consumer is ``completion_event``, or without it the MAIN work + queued so far: the MegaMoE launch that read this generation, if any. + """ self._check_owner() if ( self._plan_part is not None @@ -447,9 +429,6 @@ def finish(self, completion_event: torch.cuda.Event | None = None) -> None: or self._generation != self._finished_generation + 1 ): raise RuntimeError("finish must pair exactly once with a completed plan") - if completion_event is None: - completion_event = self._consumer_done - self._record_event(self._consumer_done_handle, self._execution_stream_handle) self.broadcaster.release_generation_after(self._generation, completion_event) self._pending_release = (self._generation, completion_event) self._finished_generation = self._generation diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py index 10f0fb48ef84..eddea7afada0 100644 --- a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py @@ -34,11 +34,6 @@ def spin_cycles_for_timeout(seconds: int | None) -> int: return seconds * ASSUMED_CLOCK_HZ -def _check_spin_cycles(spin_cycles: int) -> None: - if type(spin_cycles) is not int or spin_cycles <= 0: - raise ValueError("spin_cycles must be a positive exact int") - - def recommend_cuda_scheduler_ctas( ep_size: int, max_tokens_per_rank: int, @@ -193,33 +188,6 @@ def __init__(self, config: CudaSchedulerConfig, device: torch.device): ) -def enter_cuda_stream(stream: torch.cuda.Stream) -> tuple: - """Make ``stream`` and its device current; ``exit_cuda_stream`` restores both. - - Same effect as entering ``torch.cuda.device(stream.device)`` and then - ``torch.cuda.stream(stream)``, without their Python overhead on the - per-layer submission path. Returns the state to pass to ``exit_cuda_stream``. - """ - index = stream.device_index - prev_device = torch._C._cuda_exchangeDevice(index) - prev_stream = torch._C._cuda_getCurrentStream(index) - torch._C._cuda_setStream( - stream_id=stream.stream_id, - device_index=index, - device_type=stream.device_type, - ) - return prev_device, prev_stream - - -def exit_cuda_stream(state: tuple) -> None: - """Restore the stream and device saved by ``enter_cuda_stream``.""" - prev_device, (stream_id, device_index, device_type) = state - torch._C._cuda_setStream( - stream_id=stream_id, device_index=device_index, device_type=device_type - ) - torch._C._cuda_maybeExchangeDevice(prev_device) - - def _stream_handle(stream: object) -> int: if isinstance(stream, bool): raise TypeError("stream must be a CUDA stream handle") @@ -236,9 +204,11 @@ def _stream_handle(stream: object) -> int: class CudaPhysicalSlotScheduler: """Preallocated HALO-Q outputs and workspace for the fused scheduler + copy kernel. - The GPU-direct copy endpoint binds it once (``bind_gpu_direct``) and - launches it with the copy (``moe_rebalance_tma_submit_fused``), taking the - scheduler's arguments from ``fused_submit_args``. + The GPU-direct copy endpoint binds it once (``bind_gpu_direct``), binds + the scheduler's arguments from ``fused_launch_args`` into its native + context, and then launches it with the copy + (``moe_rebalance_tma_submit_fused``), passing the routes, the route + target and the spin bound. """ def __init__(self, config: CudaSchedulerConfig, device: str | torch.device = "cuda"): @@ -257,8 +227,8 @@ def __init__(self, config: CudaSchedulerConfig, device: str | torch.device = "cu self.outputs = CudaSchedulerOutputs(config, scheduler_device) self.out = self.outputs self.status = torch.zeros(8, dtype=torch.int32, device=scheduler_device) - # Every scheduler argument after `routes` is fixed for the endpoint's - # lifetime, so a launch passes these cached tensors and config values. + # These tensors and the config values are fixed for the endpoint's + # lifetime, so the copy binds them once; see fused_launch_args(). self.partial = torch.zeros( config.ctas * config.logical_expert_count, dtype=torch.int32, @@ -341,35 +311,24 @@ def bind_gpu_direct(self, stream: object) -> None: raise RuntimeError("EP peer bases are not connected") self._bound_stream = _stream_handle(stream) - def fused_submit_args( - self, routes: torch.Tensor, stream_handle: int, *, spin_cycles: int - ) -> tuple: - """Scheduler arguments of the bound copy's fused launch, ``routes`` - through ``valid_route_count`` of moe_rebalance_tma_submit_fused. - - ``routes`` is a caller-owned CUDA int32[T,K] tensor, read without input - staging. T may be zero and may differ across ranks, up to - max_tokens_per_rank. The kernel never reads outside those T rows and - writes -1 to the unused output tail. All ranks, including empty ranks, - must launch each generation in the same collective order. - - The caller must order input production before the bound static stream - and keep routes alive and unmodified until the scheduler completes. - Metadata checks below do not launch GPU work or synchronize; invalid - expert values are reported by the kernel through check_status(). - ``spin_cycles`` bounds this launch's cross-rank waits. + def fused_launch_args(self) -> tuple: + """Scheduler arguments that every fused launch of the bound copy + shares, ``out_slots`` through ``threads`` of + moe_rebalance_tma_bind_fused. + + A launch adds its route target, its spin bound and its routes: a + caller-owned CUDA int32[T,K] tensor, read without input staging. T may + be zero and may differ across ranks, up to max_tokens_per_rank. The + kernel never reads outside those T rows and writes -1 to the unused + output tail. All ranks, including empty + ranks, must launch each generation in the same collective order. The + caller keeps the routes alive and unmodified until the scheduler has + read them; the kernel reports invalid expert values through + check_status(). """ - _check_spin_cycles(spin_cycles) if self._bound_stream is None: raise RuntimeError("a fused launch requires a bound GPU-direct copy") - if stream_handle != self._bound_stream: - raise RuntimeError( - "CUDA scheduler uses one static stream because its workspace " - "and reusable barrier epochs are shared" - ) - valid_tokens = self._check_routes(routes) return ( - routes, self.outputs.physical_slot_ids, self.outputs.hot_expert_ids, self.outputs.hot_expert_group_level, @@ -385,25 +344,12 @@ def fused_submit_args( self.cfg.logical_expert_count, self.cfg.extra_slots_per_rank, self.cfg.local_rank, + self.cfg.topk, self.cfg.route_count, self.cfg.ctas, self.cfg.threads, - spin_cycles, - valid_tokens * self.cfg.topk, ) - def _check_routes(self, routes: torch.Tensor) -> int: - if not isinstance(routes, torch.Tensor): - raise TypeError("routes must be a torch.Tensor") - if (routes.ndim != 2 or routes.shape[1] != self.cfg.topk - or routes.shape[0] > self.cfg.max_tokens_per_rank): - raise ValueError("routes must have shape [T,topk] with 0 <= T <= max_tokens_per_rank") - if routes.dtype != torch.int32 or not routes.is_contiguous(): - raise ValueError("routes must be contiguous torch.int32") - if not routes.is_cuda or routes.device != self.device: - raise ValueError("routes must be on the scheduler CUDA device") - return routes.shape[0] - def check_status(self) -> None: """Raise if any generation faulted since the last `reset_status()`. diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py index 93c6cd0eee1c..6f4a8975fa47 100644 --- a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py @@ -14,7 +14,6 @@ from .arena import BoundLiveWeightArena, LiveWeightArena, _TensorIdentity from ._util import check_cuda -from ..cuda_scheduler.runtime import enter_cuda_stream, exit_cuda_stream from ..geometry import hierarchy_group_sizes from .geometry import BundleLayout @@ -344,10 +343,14 @@ def __init__( )) if int(torch.ops.trtllm.moe_rebalance_tma_current_gen(ctx)) != 1: raise RuntimeError("hierarchical SAMI context is not fresh") + # The per-layer launch and release ops take this handle of ctx; + # passing the custom-class object costs microseconds per call. + handle = int(torch.ops.trtllm.moe_rebalance_tma_handle(ctx)) except BaseException: torch.ops.trtllm.moe_rebalance_tma_destroy(ctx) raise self.ctx = ctx + self._handle = handle self._lock = threading.Lock() self._scheduler: object | None = None @@ -406,8 +409,15 @@ def _validate_stream_context(self, handle: int) -> None: if stream_context != current_context: raise ValueError("scheduler stream is from a different CUDA context") - def bind_scheduler(self, scheduler: object, stream: object) -> None: - """Bind the four-field CUDA scheduler and one static stream for life.""" + def bind_scheduler( + self, scheduler: object, stream: object, route_epoch: torch.Tensor + ) -> None: + """Bind the four-field CUDA scheduler and one static stream for life. + + The native context also binds every argument that the fused launches + share, including ``route_epoch``, the caller-owned int64 that each + launch raises once its routes are written. + """ with self._lock: if self._handoff_aborted: @@ -484,7 +494,7 @@ def bind_scheduler(self, scheduler: object, stream: object) -> None: if stream_device != scheduler_device: raise ValueError("scheduler stream and output devices differ") bind = getattr(scheduler, "bind_gpu_direct", None) - fused_args = getattr(scheduler, "fused_submit_args", None) + fused_args = getattr(scheduler, "fused_launch_args", None) if not callable(bind) or not callable(fused_args): raise TypeError("scheduler lacks the fused copy launch lifecycle") workspace = getattr(scheduler, "plan_workspace", None) @@ -515,12 +525,16 @@ def bind_scheduler(self, scheduler: object, stream: object) -> None: capacity, self.helper_residency, ) + torch.ops.trtllm.moe_rebalance_tma_bind_fused( + self.ctx, self.flags_mc, route_epoch, *fused_args(), handle + ) except BaseException: self._handoff_aborted = True raise self._scheduler = scheduler self._outputs = outputs self._stream = handle + # The native context launches on this stream by handle; keep it alive. self._stream_owner = stream def _validate_reuse_authority(self, authority: object) -> None: @@ -568,20 +582,24 @@ def bind_generation_reuse_authority(self, authority: object) -> None: self._reuse_authority = authority def release_generation_after( - self, generation: int, completion_event: object + self, generation: int, completion_event: object | None = None ) -> None: - """Order the copy stream after this rank's live-bank consumer event. + """Order the copy stream after this rank's live-bank consumer. - This establishes only local stream ordering. The bound reuse authority - must separately attest its bounded all-rank barrier before another - generation may overwrite the single live bank. + The consumer is ``completion_event``, or without it the work already + queued on the current stream. This establishes only local stream + ordering. The bound reuse authority must separately attest its + bounded all-rank barrier before another generation may overwrite the + single live bank. """ if type(generation) is not int or generation <= 0: raise ValueError("generation must be a positive exact int") - event_handle = _event_handle(completion_event) - if event_handle <= 0: - raise ValueError("completion_event must expose a positive CUDA event") + event_handle = 0 + if completion_event is not None: + event_handle = _event_handle(completion_event) + if event_handle <= 0: + raise ValueError("completion_event must expose a positive CUDA event") with self._lock: if self._handoff_aborted or self._scheduler is None: raise RuntimeError("hierarchical scheduler handoff is not operational") @@ -589,10 +607,13 @@ def release_generation_after( raise ValueError("generation is not the active live bank") if self._local_release_generation >= generation: raise RuntimeError("generation already has a local release event") - check_cuda( - self.cuda.cuStreamWaitEvent(self._stream, event_handle, 0), - "cuStreamWaitEvent(live-bank consumer)", - ) + if completion_event is None: + torch.ops.trtllm.moe_rebalance_tma_release(self._handle) + else: + check_cuda( + self.cuda.cuStreamWaitEvent(self._stream, event_handle, 0), + "cuStreamWaitEvent(live-bank consumer)", + ) self._local_release_generation = generation def mark_collective_reuse_safe( @@ -633,17 +654,18 @@ def submit_fused( self, outputs: object, routes: torch.Tensor, - route_epoch: torch.Tensor, route_target: int, *, spin_cycles: int, ) -> HierarchicalCopyTicket: """Enqueue the bound scheduler and this GPU-direct copy as one kernel. - Each scheduler CTA adds one to ``route_epoch`` once its routes are - written, which makes it reach ``route_target``; the copy and its READY - publication follow in the same CTAs. ``spin_cycles`` bounds the - launch's cross-rank waits and its wait for the route epoch. + The kernel runs on the bound stream after the work already queued on + the current stream, which must produce ``routes``. Each scheduler CTA + adds one to the bound route epoch once its routes are written, which + makes it reach ``route_target``; the copy and its READY publication + follow in the same CTAs. ``spin_cycles`` bounds the launch's + cross-rank waits and its wait for the route epoch. """ with self._lock: @@ -655,16 +677,9 @@ def submit_fused( raise RuntimeError("outputs are not from the bound CUDA scheduler") self._require_live_bank_reusable() try: - scheduler_args = self._scheduler.fused_submit_args( - routes, self._stream, spin_cycles=spin_cycles + torch.ops.trtllm.moe_rebalance_tma_submit_fused( + self._handle, routes, route_target, spin_cycles ) - state = enter_cuda_stream(self._stream_owner) - try: - torch.ops.trtllm.moe_rebalance_tma_submit_fused( - self.ctx, self.flags_mc, route_epoch, route_target, *scheduler_args - ) - finally: - exit_cuda_stream(state) self._active_generation = self._generation except BaseException: self._handoff_aborted = True @@ -682,6 +697,7 @@ def close(self) -> None: with self._lock, torch.cuda.device(self.device): if self.ctx is None: return + self._handle = 0 torch.ops.trtllm.moe_rebalance_tma_destroy(self.ctx) self.ctx = None self._scheduler = None diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index e85edc2b43c5..46dc62d845f9 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -231,44 +231,53 @@ def build_layer(self, comm): def test_hierarchical_fused_submit_launches_bound_scheduler_and_copy(): source = _TORCH_ROOT / "moe/kernels/megamoe_scheduler/sami/hierarchical.py" - launches = [] + native = [] - def submit_fused(ctx, flags_mc, route_epoch, route_target, *scheduler_args): - launches.append((ctx, flags_mc, route_epoch, route_target, scheduler_args)) + def submit_fused(handle, routes, route_target, spin_cycles): + native.append(("launch", handle, routes, route_target, spin_cycles)) + + def release(handle): + native.append(("release", handle)) torch = SimpleNamespace( - ops=SimpleNamespace(trtllm=SimpleNamespace(moe_rebalance_tma_submit_fused=submit_fused)), + ops=SimpleNamespace( + trtllm=SimpleNamespace( + moe_rebalance_tma_submit_fused=submit_fused, + moe_rebalance_tma_release=release, + ) + ), ) - stream_switches = [] ticket_cls = _extract_class( source, "HierarchicalCopyTicket", None, {"dataclass": dataclasses.dataclass, "field": dataclasses.field}, ) + event_handle = _extract_function(source, "_event_handle", {}) broadcast_cls = _extract_class( source, "HierarchicalSamiWeightBroadcast", - {"submit_fused", "_require_live_bank_reusable", "mark_collective_reuse_safe"}, + { + "submit_fused", + "release_generation_after", + "_require_live_bank_reusable", + "mark_collective_reuse_safe", + }, { "torch": torch, "HierarchicalCopyTicket": ticket_cls, - "enter_cuda_stream": lambda stream: stream_switches.append(("enter", stream)) or stream, - "exit_cuda_stream": lambda state: stream_switches.append(("exit", state)), + "_event_handle": event_handle, + "check_cuda": lambda result, operation: result, }, ) outputs = object() - requests = [] - - def fused_submit_args(routes, stream_handle, *, spin_cycles): - requests.append((routes, stream_handle, spin_cycles)) - return ("scheduler", routes) broadcast = broadcast_cls.__new__(broadcast_cls) broadcast._lock = nullcontext() broadcast._handoff_aborted = False broadcast._outputs = outputs - broadcast._scheduler = SimpleNamespace(fused_submit_args=fused_submit_args) + # The native context binds every scheduler argument; a launch passes only its routes. + broadcast._scheduler = SimpleNamespace() authority = object() broadcast._reuse_authority = authority @@ -279,41 +288,55 @@ def revalidate(_authority): broadcast._validate_reuse_authority = revalidate broadcast._active_generation = 0 broadcast._collective_safe_generation = 0 + broadcast._local_release_generation = 0 broadcast._generation = 1 broadcast._stream = 19 - broadcast._stream_owner = object() broadcast.device = 0 broadcast.ctx = "ctx" - broadcast.flags_mc = 4096 + # The per-layer ops take the context's handle, not the custom-class object. + broadcast._handle = 4242 broadcast.terminal_flags = "terminals" + broadcast.cuda = SimpleNamespace( + cuStreamWaitEvent=lambda stream, event, flags: native.append(("wait", stream, event)) + ) - ticket = broadcast.submit_fused(outputs, "routes", "epoch", 8, spin_cycles=12345) + ticket = broadcast.submit_fused(outputs, "routes", 8, spin_cycles=12345) - assert requests == [("routes", 19, 12345)] - assert launches == [("ctx", 4096, "epoch", 8, ("scheduler", "routes"))] - # The launch runs on the copy stream and restores the caller's stream afterwards. - owner = broadcast._stream_owner - assert stream_switches == [("enter", owner), ("exit", owner)] + # One native call orders the launch after the current stream and enqueues + # it on the bound copy stream; Python switches no stream. + assert native == [("launch", 4242, "routes", 8, 12345)] assert (ticket.generation, ticket.command_count) == (1, 1) assert (broadcast._active_generation, broadcast._generation) == (1, 2) # The bank stays leased until the consumer of generation 1 is proven done. with pytest.raises(RuntimeError, match="still leased"): - broadcast.submit_fused(outputs, "routes", "epoch", 16, spin_cycles=12345) + broadcast.submit_fused(outputs, "routes", 16, spin_cycles=12345) with pytest.raises(RuntimeError, match="bound CUDA scheduler"): - broadcast.submit_fused(object(), "routes", "epoch", 16, spin_cycles=12345) + broadcast.submit_fused(object(), "routes", 16, spin_cycles=12345) # Only the bound authority, after the local consumer release, frees the bank. with pytest.raises(PermissionError, match="identity mismatch"): broadcast.mark_collective_reuse_safe(object(), 1) - broadcast._local_release_generation = 0 with pytest.raises(RuntimeError, match="local consumer release"): broadcast.mark_collective_reuse_safe(authority, 1) - broadcast._local_release_generation = 1 + native.clear() + # Without an event, the consumer is the work queued on the current stream, + # which one native call orders the copy stream after. + broadcast.release_generation_after(1) + assert native == [("release", 4242)] + with pytest.raises(RuntimeError, match="already has a local release"): + broadcast.release_generation_after(1) broadcast.mark_collective_reuse_safe(authority, 1) assert broadcast._collective_safe_generation == 1 - ticket = broadcast.submit_fused(outputs, "routes", "epoch", 16, spin_cycles=12345) + ticket = broadcast.submit_fused(outputs, "routes", 16, spin_cycles=12345) assert ticket.generation == 2 + # A caller-recorded consumer event is waited for on the copy stream instead. + native.clear() + broadcast.release_generation_after(2, SimpleNamespace(cuda_event=23)) + assert native == [("wait", 19, 23)] + with pytest.raises(ValueError, match="positive CUDA event"): + broadcast.release_generation_after(3, SimpleNamespace(cuda_event=0)) + @pytest.mark.parametrize( ( @@ -1245,11 +1268,6 @@ def test_fused_submit_orders_generations_and_defers_the_route_wait(auxiliary_sms group.auxiliary_sms = auxiliary_sms group.copy_stream = group.scheduler_stream = copy group._stream_handle = 19 - group._input_ready_handle = 1 - group._consumer_done_handle = 3 - group._consumer_done = SimpleNamespace(cuda_event=3) - route_epoch = object() - group._route_epoch = route_epoch group._route_epoch_ptr = 77 # Generation 1 launches during warmup, generation 2 while serving. bounds = iter((_WARMUP_SPIN_CYCLES, _SERVING_SPIN_CYCLES)) @@ -1259,11 +1277,11 @@ def driver_call(*entry): trace.append(entry) return (0,) + # The planner records no event itself: the native launch and release + # order the copy stream after MAIN. group._driver = SimpleNamespace( CUresult=SimpleNamespace(CUDA_SUCCESS=0), CUstreamWaitValue_flags=SimpleNamespace(CU_STREAM_WAIT_VALUE_GEQ="geq"), - cuEventRecord=lambda event, stream: driver_call("record", event, stream), - cuStreamWaitEvent=lambda stream, event, flags: driver_call("wait", stream, event), cuStreamWaitValue64=lambda stream, address, value, flags: driver_call( "wait_value", stream, address, value, flags ), @@ -1272,8 +1290,8 @@ def driver_call(*entry): launch_bounds = [] - def submit_fused(result, routes, epoch, target, *, spin_cycles): - assert result is outputs and epoch is route_epoch + def submit_fused(result, routes, target, *, spin_cycles): + assert result is outputs generation = group._scheduled_generation trace.append(("FUSED", routes.shape[0], target, generation)) launch_bounds.append(spin_cycles) @@ -1283,7 +1301,7 @@ def submit_fused(result, routes, epoch, target, *, spin_cycles): group.broadcaster = SimpleNamespace( submit_fused=submit_fused, release_generation_after=lambda generation, event: trace.append( - ("consumer_release", generation, event.cuda_event) + ("consumer_release", generation, event) ), mark_collective_reuse_safe=lambda provider, generation: trace.append(("lease", generation)), ) @@ -1308,9 +1326,9 @@ def submit_fused(result, routes, epoch, target, *, spin_cycles): group.finish() assert trace.count(route_wait) == 1 assert not any(entry[0] == "lifetime" for entry in trace) - assert trace.index(("wait", 19, 1)) < trace.index(("FUSED", rows, target, generation)) assert trace.index(("FUSED", rows, target, generation)) < trace.index(route_wait) - assert trace[-2:] == [("record", 3, 44), ("consumer_release", generation, 3)] + # Without an event, the release follows the MAIN work queued so far. + assert trace[-1] == ("consumer_release", generation, None) if generation > 1: # The previous consumer is proven reusable before the next launch. assert trace.index(("lease", generation - 1)) < trace.index( @@ -1333,17 +1351,8 @@ def submit_fused(result, routes, epoch, target, *, spin_cycles): group.plan_schedule(_Routes(8, trace, dtype="int64")) assert trace == [] - group._driver.cuStreamWaitEvent = lambda *args: (17,) - with pytest.raises(RuntimeError, match="cuStreamWaitEvent"): - group.plan_schedule(_Routes(8, trace)) - assert not any(entry[0] == "FUSED" for entry in trace) - assert group._plan_part is None - # Nothing was launched, so the routes were never handed to HALO-Q. - assert group._routes_in_flight is None - - # A submission that fails after the launch keeps its routes alive, and the - # group refuses further plans rather than let the allocator reuse them. - group._driver.cuStreamWaitEvent = lambda *args: (0,) + # A submission that fails may have launched, so it keeps its routes alive, + # and the group refuses further plans rather than let the allocator reuse them. group._pending_release = None group._spin_bound = SimpleNamespace(cycles=lambda: _SERVING_SPIN_CYCLES) group.broadcaster.submit_fused = lambda *args, **kwargs: SimpleNamespace(generation=99) @@ -1435,7 +1444,6 @@ def _module_constants(path: Path, names: set[str]) -> dict: def _scheduler_spin_namespace() -> dict: namespace = _module_constants(_SCHEDULER_RUNTIME, {"SPIN_CYCLES", "ASSUMED_CLOCK_HZ"}) _extract_function(_SCHEDULER_RUNTIME, "spin_cycles_for_timeout", namespace) - _extract_function(_SCHEDULER_RUNTIME, "_check_spin_cycles", namespace) return namespace @@ -1470,9 +1478,6 @@ def test_scheduler_spin_cycles_convert_timeouts_at_the_nominal_clock(): for seconds in (0, -1, 2.5, True): with pytest.raises(ValueError, match="positive exact int"): to_cycles(seconds) - for cycles in (0, -1, 1.0, True): - with pytest.raises(ValueError, match="positive exact int"): - namespace["_check_spin_cycles"](cycles) @pytest.mark.parametrize( @@ -1522,22 +1527,35 @@ def test_scheduler_spin_bound_follows_warmup_transitions_and_autotuning( assert bound.cycles() == warmup_cycles -def test_fused_launch_passes_its_spin_bound_to_the_kernel(): - # Map the scheduler arguments onto the registered operator schema rather than hard-coding them. +def _tma_op_parameters(name: str) -> List[str]: + """Parameter names of a registered moe_rebalance_tma operator schema.""" op_source = (_ROOT / "cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp").read_text() schemas = ["".join(re.findall(r'"([^"]*)"', block)) for block in op_source.split("m.def(")[1:]] - schema = next(text for text in schemas if text.startswith("moe_rebalance_tma_submit_fused(")) - parameters = re.fullmatch(r"moe_rebalance_tma_submit_fused\((.*)\) -> \(\)", schema).group(1) - names = [parameter.split()[-1] for parameter in parameters.split(",")] - scheduler_names = names[names.index("routes") :] + schema = next(text for text in schemas if text.startswith(f"{name}(")) + parameters = re.fullmatch(rf"{name}\((.*)\) -> \(\)", schema).group(1) + return [parameter.split()[-1] for parameter in parameters.split(",")] + + +def test_fused_launch_binds_shared_arguments_and_passes_only_routes(): + # Map the scheduler arguments onto the registered operator schemas rather than hard-coding them. + bind_names = _tma_op_parameters("moe_rebalance_tma_bind_fused") + scheduler_names = bind_names[bind_names.index("out_slots") : bind_names.index("copy_stream")] + # A launch passes its routes, route target and spin bound; the bind holds the rest. + # They take the context's handle, since a custom-class argument costs + # microseconds of dispatch per call. + assert _tma_op_parameters("moe_rebalance_tma_submit_fused") == [ + "handle", + "routes", + "route_target", + "spin_cycles", + ] + assert _tma_op_parameters("moe_rebalance_tma_release") == ["handle"] - namespace = _scheduler_spin_namespace() - namespace["torch"] = SimpleNamespace(Tensor=_Routes, int32="int32") scheduler_cls = _extract_class( _SCHEDULER_RUNTIME, "CudaPhysicalSlotScheduler", - {"fused_submit_args", "_check_routes"}, - namespace, + {"fused_launch_args"}, + {}, ) scheduler = scheduler_cls.__new__(scheduler_cls) scheduler.cfg = SimpleNamespace( @@ -1567,61 +1585,82 @@ def test_fused_launch_passes_its_spin_bound_to_the_kernel(): "route_prefix", ): setattr(scheduler, name, name) - scheduler.device = "cuda:0" + scheduler._bound_stream = None + with pytest.raises(RuntimeError, match="bound GPU-direct copy"): + scheduler.fused_launch_args() scheduler._bound_stream = 19 - routes = _Routes(16, []) - routes.is_cuda = True - args = scheduler.fused_submit_args(routes, 19, spin_cycles=_WARMUP_SPIN_CYCLES) + args = scheduler.fused_launch_args() assert len(args) == len(scheduler_names) - assert args[scheduler_names.index("spin_cycles")] == _WARMUP_SPIN_CYCLES - assert args[scheduler_names.index("valid_route_count")] == 16 * 6 - with pytest.raises(ValueError, match="positive exact int"): - scheduler.fused_submit_args(routes, 19, spin_cycles=0) - # Every fused launch shares the workspace, so it must use the bound stream. - with pytest.raises(RuntimeError, match="one static stream"): - scheduler.fused_submit_args(routes, 20, spin_cycles=_WARMUP_SPIN_CYCLES) - - -def test_raw_stream_switch_restores_the_previous_stream_and_device(): - calls = [] - current = {"device": 0, "streams": {0: (10, 0, 1), 3: (30, 3, 1)}} - - def exchange_device(index): - calls.append(("exchange", index)) - previous, current["device"] = current["device"], index - return previous - - def set_stream(stream_id, device_index, device_type): - calls.append(("set", stream_id, device_index)) - current["streams"][device_index] = (stream_id, device_index, device_type) - - def maybe_exchange_device(index): - calls.append(("restore_device", index)) - current["device"] = index - return index - - raw = SimpleNamespace( - _cuda_exchangeDevice=exchange_device, - _cuda_getCurrentStream=lambda index: current["streams"][index], - _cuda_setStream=set_stream, - _cuda_maybeExchangeDevice=maybe_exchange_device, - ) - source = _TORCH_ROOT / "moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py" - namespace = {"torch": SimpleNamespace(_C=raw, cuda=SimpleNamespace(Stream=object))} - enter = _extract_function(source, "enter_cuda_stream", namespace) - leave = _extract_function(source, "exit_cuda_stream", namespace) - - state = enter(SimpleNamespace(device_index=3, stream_id=77, device_type=1)) - assert current["device"] == 3 and current["streams"][3] == (77, 3, 1) - leave(state) - assert current == {"device": 0, "streams": {0: (10, 0, 1), 3: (30, 3, 1)}} - assert calls == [ - ("exchange", 3), - ("set", 77, 3), - ("set", 30, 3), - ("restore_device", 0), - ] + expected = { + "out_slots": "slots", + "out_ids": "ids", + "out_levels": "levels", + "out_owners": "owners", + "peer_bases": "peer_base", + "route_capacity": 8192 * 6, + "topk": 6, + "ctas": 4, + "threads": 512, + "ep": 8, + "experts": 384, + "helpers": 4, + "local_rank": 0, + } + for name, value in expected.items(): + assert args[scheduler_names.index(name)] == value + + +def _cpp_method_body(source: str, signature: str) -> str: + """Body of the C++ function whose definition starts with ``signature``.""" + open_brace = source.index("{", source.index(signature)) + depth = 0 + for index in range(open_brace, len(source)): + if source[index] == "{": + depth += 1 + elif source[index] == "}": + depth -= 1 + if depth == 0: + return source[open_brace + 1 : index] + raise AssertionError(f"unbalanced body of {signature}") + + +def _assert_in_order(body: str, calls: Tuple[str, ...]) -> None: + positions = [body.index(call) for call in calls] + assert positions == sorted(positions), calls + + +def test_native_launch_and_release_order_the_copy_stream_after_main(): + # The cross-stream ordering lives in the native context, so pin its call order: + # HALO-Q must not read routes before MAIN wrote them, and the next generation's + # copy must not overwrite the live bank before MAIN's consumer of it. + source = (_ROOT / "cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp").read_text() + submit = _cpp_method_body(source, "void submitFused(torch::Tensor const& routes") + _assert_in_order( + submit, + ( + "TORCH_CHECK_VALUE(routes.is_cuda()", + "TORCH_CHECK_VALUE(spinCycles > 0", + "cudaStream_t const execution = currentExecutionStream();", + "cudaEventRecord(mInputReady, execution)", + "cudaStreamWaitEvent(mCopyStream, mInputReady, 0)", + "megamoe_tma_copy_submit_fused(", + ), + ) + launch = submit[submit.index("megamoe_tma_copy_submit_fused(") :] + assert "reinterpret_cast(mCopyStream)" in launch + release = _cpp_method_body(source, "void releaseAfterCurrentStream()") + _assert_in_order( + release, + ( + "cudaStream_t const execution = currentExecutionStream();", + "cudaEventRecord(mConsumerDone, execution)", + "cudaStreamWaitEvent(mCopyStream, mConsumerDone, 0)", + ), + ) + # The stream a launch or release follows can never be the copy stream itself. + current = _cpp_method_body(source, "cudaStream_t currentExecutionStream() const") + assert "TORCH_CHECK(stream != mCopyStream" in current def test_megamoe_cute_weight_cache_keys_every_view_field(): diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index 4e4b79cc69de..ca01bc542577 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -2660,7 +2660,6 @@ def test_megamoe_cutedsl_cache_derived_state_reuses_rebalance_binding( @pytest.mark.cpu_only def test_rebalance_discard_plan_orders_release_and_allows_next_reuse(): trace = [] - completion_event = object() outputs = SimpleNamespace(physical_slot_ids=[3, 4]) part = (outputs, 2) copy_stream = object() @@ -2674,8 +2673,6 @@ def test_rebalance_discard_plan_orders_release_and_allows_next_reuse(): group._route_wait_generation = 0 group._pending_release = None group.plan_calls = 0 - group._consumer_done = completion_event - group._consumer_done_handle = 17 group._execution_stream_handle = 23 group.copy_stream = copy_stream group.scheduler_stream = copy_stream @@ -2684,20 +2681,18 @@ def wait_for_routes(): trace.append("wait_routes") group._route_wait_generation = group._generation - def record_event(event_handle, stream_handle): - assert (event_handle, stream_handle) == (17, 23) - trace.append("record_consumer_done") + def release_generation_after(generation, event): + # Without an event, the release follows the MAIN work queued so far. + assert (generation, event) == (1, None) + trace.append("release_generation") broadcaster = SimpleNamespace( - release_generation_after=MagicMock( - side_effect=lambda generation, event: trace.append("release_generation") - ), + release_generation_after=MagicMock(side_effect=release_generation_after), mark_collective_reuse_safe=MagicMock( side_effect=lambda authority, generation: trace.append("prove_reuse") ), ) group.wait_for_routes = wait_for_routes - group._record_event = record_event group.broadcaster = broadcaster group.discard_plan(part) @@ -2705,8 +2700,8 @@ def record_event(event_handle, stream_handle): assert group._plan_part is None assert group.plan_calls == 1 assert group._finished_generation == 1 - assert group._pending_release == (1, completion_event) - assert trace == ["wait_routes", "record_consumer_done", "release_generation"] + assert group._pending_release == (1, None) + assert trace == ["wait_routes", "release_generation"] # The next schedule advances this counter before proving that generation 1 # can be reused. Exercise the real lease check to show cleanup left a valid From 305252efb6f6b2a93de4130d02c6be41dffb7e24 Mon Sep 17 00:00:00 2001 From: jintaop Date: Sun, 25 Oct 2026 06:46:44 -0700 Subject: [PATCH 39/59] [None][perf] Validate shared FC12 weights once per weight set Each shared-expert FC12 call re-validated its four weights, rebuilt their layout signature, built a torch.cuda.Stream, entered a device guard and let the AutoTuner stringify a unique_id that embeds that signature. On a GB200 host this wrapper took about 107 us per call, as much as or more than the kernel at small token counts. Validate each weight set once, keyed by tensor identity, read the raw current stream, take the device guard only when a runner is created, and build the unique_id string once. The AutoTuner cache keys are unchanged. The runner also reuses its token-count views and returns a view of its output. The same GB200 call now takes about 56 us. Signed-off-by: jintaop --- .../_torch/modules/megamoe_shared_mlp.py | 4 +- .../_torch/moe/kernels/megamoe_shared_fc12.py | 150 ++++++++++++------ .../moe/test_megamoe_shared_fc12_autotune.py | 60 ++++++- 3 files changed, 164 insertions(+), 50 deletions(-) diff --git a/tensorrt_llm/_torch/modules/megamoe_shared_mlp.py b/tensorrt_llm/_torch/modules/megamoe_shared_mlp.py index 0e0d416c437e..798306e7415a 100644 --- a/tensorrt_llm/_torch/modules/megamoe_shared_mlp.py +++ b/tensorrt_llm/_torch/modules/megamoe_shared_mlp.py @@ -136,6 +136,7 @@ def __init__( if self._fc12_max_tokens <= 0: raise ValueError("Shared FC12 max_num_tokens must be positive") self._fc12_source_state: tuple[tuple[int, int, int | None], ...] | None = None + self._fc12_sm_count: int | None = None for name in ("fc1_weight", "fc1_weight_sf", "fc2_weight", "fc2_weight_sf"): self.register_buffer(f"_fc12_{name}", None, persistent=False) @@ -220,6 +221,7 @@ def prepare_fc12_weights(self) -> None: self._fc12_fc1_weight_sf = packed_fc1_sf self._fc12_fc2_weight = fc2.unsqueeze(0) self._fc12_fc2_weight_sf = fc2_sf.unsqueeze(0) + self._fc12_sm_count = self.sm_count self._fc12_source_state = source_state def forward( @@ -252,6 +254,6 @@ def forward( self._fc12_fc2_weight, self._fc12_fc2_weight_sf, swiglu_limit=self.swiglu_limit, - sm_count=self.sm_count, + sm_count=self._fc12_sm_count, max_tokens=self._fc12_max_tokens, ) diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_shared_fc12.py b/tensorrt_llm/_torch/moe/kernels/megamoe_shared_fc12.py index 0106651e7126..a7b7e477a972 100644 --- a/tensorrt_llm/_torch/moe/kernels/megamoe_shared_fc12.py +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_shared_fc12.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """Single-expert use of MegaMoE's complete fused FC1/activation/FC2 kernel.""" +import weakref from itertools import product from threading import Condition, RLock @@ -164,6 +165,7 @@ def __init__( self._cuda_stream = cuda.CUstream(stream_handle) # A slice supplies the current E=1 token count without a per-call GPU write. self.token_counts = torch.arange(max_tokens + 1, dtype=torch.int32, device=device) + self._token_count_views: dict[int, torch.Tensor] = {} token_rows = cute.sym_int64(divisibility=1) scale_rows = cute.sym_int64(divisibility=128) @@ -230,6 +232,11 @@ def __call__( fc2_weight_sf: torch.Tensor, ) -> torch.Tensor: token_rows = activation.shape[0] + token_count = self._token_count_views.get(token_rows) + if token_count is None: + token_count = self._token_count_views[token_rows] = self.token_counts[ + token_rows : token_rows + 1 + ] quantized, scales = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0.sm_budget( activation, reserved_sms=self.reserved_sms ) @@ -252,14 +259,14 @@ def __call__( fc2_output=output, local_workspace=self.workspace_pointer, stream=self.stream_handle, - expert_token_sizes=self.token_counts[token_rows : token_rows + 1], + expert_token_sizes=token_count, expert_token_prefix_sum=None, topk_scores=None, fc1_alpha=None, fc2_alpha=None, fc1_norm_const=None, ) - return output[:, 0, :] + return output.view(token_rows, self.hidden_size) _SHARED_FC12_CACHE_LOCK = RLock() @@ -409,19 +416,23 @@ def __init__( swiglu_limit, stream_handle, ) - self._unique_id = ( - _CATALOG_VERSION, - device_index, - properties.major, - properties.minor, - properties.multi_processor_count, - sm_count, - hidden_size, - intermediate_size, - max_tokens, - swiglu_limit, - "K128_input_K32_intermediate_R128c4_gateup16", - input_signature, + # The AutoTuner keys its cache by str(unique_id()) on every dispatch, + # so build that string once. + self._unique_id = str( + ( + _CATALOG_VERSION, + device_index, + properties.major, + properties.minor, + properties.multi_processor_count, + sm_count, + hidden_size, + intermediate_size, + max_tokens, + swiglu_limit, + "K128_input_K32_intermediate_R128c4_gateup16", + input_signature, + ) ) self.buckets = _token_buckets(max_tokens) self.tuning_config = TuningConfig( @@ -571,6 +582,64 @@ def matches(key: tuple) -> bool: _end_shared_fc12_release(device_index, stream_key) +# Validated weight layouts by the identity of the four weight tensors. A +# layer's weights stay fixed between loads, so each set is checked once. +_WEIGHT_LAYOUTS: dict[tuple[int, int, int, int], tuple] = {} + + +def _weight_layout(weights: tuple, hidden_size: int) -> tuple[int, torch.device, tuple]: + """Validate FC12 weights once per weight set. + + Returns the intermediate size, the weights' device, and the input + signature of a contiguous BF16 ``[T, hidden_size]`` activation with them. + """ + key = (id(weights[0]), id(weights[1]), id(weights[2]), id(weights[3])) + entry = _WEIGHT_LAYOUTS.get(key) + if entry is not None: + refs, layout_hidden_size, intermediate_size, device, signature = entry + if ( + layout_hidden_size == hidden_size + and refs[0]() is weights[0] + and refs[1]() is weights[1] + and refs[2]() is weights[2] + and refs[3]() is weights[3] + ): + return intermediate_size, device, signature + fc1_weight, fc1_weight_sf, fc2_weight, fc2_weight_sf = weights + if fc2_weight.ndim != 3 or fc2_weight.shape[:2] != (1, hidden_size): + raise ValueError("shared FC12 weight must have shape [1, H, I]") + intermediate_size = fc2_weight.shape[2] + if hidden_size % 128 or intermediate_size % 128: + raise ValueError("shared FC12 hidden and intermediate dimensions must be multiples of 128") + if fc1_weight.shape != (1, 2 * intermediate_size, hidden_size): + raise ValueError("shared FC1 weight must have shape [1, 2*I, H]") + for weight in (fc1_weight, fc2_weight): + if weight.dtype != torch.float8_e4m3fn or not weight.is_contiguous(): + raise ValueError("shared FC12 weights must be contiguous E4M3 tensors") + for scale in (fc1_weight_sf, fc2_weight_sf): + if scale.dtype != torch.uint8 or not scale.is_contiguous(): + raise ValueError("shared FC12 scales must be contiguous uint8 tensors") + device = fc1_weight.device + if any(t.device != device for t in weights): + raise ValueError("shared FC12 inputs and weights must be on the same CUDA device") + signature = ( + (str(torch.bfloat16), (None, hidden_size), (hidden_size, 1)), + *((str(t.dtype), tuple(t.shape), tuple(t.stride())) for t in weights), + ) + with _SHARED_FC12_CACHE_LOCK: + stale = [k for k, e in _WEIGHT_LAYOUTS.items() if any(ref() is None for ref in e[0])] + for k in stale: + del _WEIGHT_LAYOUTS[k] + _WEIGHT_LAYOUTS[key] = ( + tuple(weakref.ref(t) for t in weights), + hidden_size, + intermediate_size, + device, + signature, + ) + return intermediate_size, device, signature + + def run_shared_fc12( activation: torch.Tensor, fc1_weight: torch.Tensor, @@ -594,51 +663,36 @@ def run_shared_fc12( raise ValueError("shared FC12 requires CUDA BF16 activations") if activation.ndim != 2 or not activation.is_contiguous(): raise ValueError("shared FC12 requires contiguous [T, H] activations") - if max_tokens <= 0 or activation.shape[0] > max_tokens: + tokens, hidden_size = activation.shape + if max_tokens <= 0 or tokens > max_tokens: raise ValueError("shared FC12 token count exceeds its configured capacity") - hidden_size = activation.shape[1] - if fc2_weight.ndim != 3 or fc2_weight.shape[:2] != (1, hidden_size): - raise ValueError("shared FC12 weight must have shape [1, H, I]") - intermediate_size = fc2_weight.shape[2] - if hidden_size % 128 or intermediate_size % 128: - raise ValueError("shared FC12 hidden and intermediate dimensions must be multiples of 128") - if fc1_weight.shape != (1, 2 * intermediate_size, hidden_size): - raise ValueError("shared FC1 weight must have shape [1, 2*I, H]") - for weight in (fc1_weight, fc2_weight): - if weight.dtype != torch.float8_e4m3fn or not weight.is_contiguous(): - raise ValueError("shared FC12 weights must be contiguous E4M3 tensors") - for scale in (fc1_weight_sf, fc2_weight_sf): - if scale.dtype != torch.uint8 or not scale.is_contiguous(): - raise ValueError("shared FC12 scales must be contiguous uint8 tensors") - if any( - t.device != activation.device - for t in (fc1_weight, fc1_weight_sf, fc2_weight, fc2_weight_sf) - ): + weights = (fc1_weight, fc1_weight_sf, fc2_weight, fc2_weight_sf) + intermediate_size, device, input_signature = _weight_layout(weights, hidden_size) + if activation.device != device: raise ValueError("shared FC12 inputs and weights must be on the same CUDA device") - if activation.shape[0] == 0: + if tokens == 0: return torch.empty_like(activation) - stream = torch.cuda.current_stream(activation.device) - inputs = [activation, fc1_weight, fc1_weight_sf, fc2_weight, fc2_weight_sf] - input_signature = tuple( - (str(t.dtype), (None, hidden_size) if index == 0 else tuple(t.shape), tuple(t.stride())) - for index, t in enumerate(inputs) - ) + device_index = device.index + stream_handle = torch._C._cuda_getCurrentRawStream(device_index) + inputs = [activation, *weights] runner_key = ( - activation.device.index, + device_index, hidden_size, intermediate_size, max_tokens, sm_count, swiglu_limit, - stream.cuda_stream, + stream_handle, input_signature, ) - active_key = _begin_shared_fc12_call(activation.device.index, stream.cuda_stream) + active_key = _begin_shared_fc12_call(device_index, stream_handle) try: - with _SHARED_FC12_CACHE_LOCK, torch.cuda.device(activation.device): - runner = _tunable_runners.get(runner_key) - if runner is None: - runner = _tunable_runners[runner_key] = SharedFc12TunableRunner(*runner_key) + runner = _tunable_runners.get(runner_key) + if runner is None: + with _SHARED_FC12_CACHE_LOCK, torch.cuda.device(device_index): + runner = _tunable_runners.get(runner_key) + if runner is None: + runner = _tunable_runners[runner_key] = SharedFc12TunableRunner(*runner_key) tuner = AutoTuner.get() selected_runner, tactic = tuner.choose_one( _SHARED_FC12_OP, [runner], runner.tuning_config, inputs diff --git a/tests/unittest/_torch/moe/test_megamoe_shared_fc12_autotune.py b/tests/unittest/_torch/moe/test_megamoe_shared_fc12_autotune.py index 1d5b8540b45e..14a0535444e7 100644 --- a/tests/unittest/_torch/moe/test_megamoe_shared_fc12_autotune.py +++ b/tests/unittest/_torch/moe/test_megamoe_shared_fc12_autotune.py @@ -8,6 +8,7 @@ import inspect import os import unittest +import weakref from contextlib import nullcontext from copy import deepcopy from itertools import product @@ -107,12 +108,12 @@ def external_stream(stream_handle, *, device): cuda=SimpleNamespace( get_device_properties=lambda device: properties, device=lambda device: nullcontext(), - current_stream=lambda device: SimpleNamespace(cuda_stream=17), ExternalStream=external_stream, synchronize=lambda device: self.device_syncs.append(device), empty_cache=lambda: setattr(self, "empty_cache_calls", self.empty_cache_calls + 1), ), empty_like=lambda x: Tensor(x.shape, x.dtype), + _C=SimpleNamespace(_cuda_getCurrentRawStream=lambda device_index: 17), ) self.ns = dict( torch=torch, @@ -124,6 +125,7 @@ def external_stream(stream_handle, *, device): RLock=RLock, deepcopy=deepcopy, product=product, + weakref=weakref, ) module = ast.parse(_SOURCE.read_text()) module.body = [ @@ -336,6 +338,62 @@ def choose(op, runners, config, tensors): self.assertEqual([output.shape[0] for output in outputs], [129, 200]) self.assertEqual([tokens for _, tokens in self.launches], [129, 200]) + def test_weight_layout_is_validated_once_per_weight_set(self): + checks = [] + + class CountedWeight(Tensor): + @property + def ndim(self): + checks.append(id(self)) + return len(self.shape) + + @ndim.setter + def ndim(self, value): + pass + + self.tuner.choose_one = lambda op, runners, config, tensors: (runners[0], -1) + activation, fc1, fc1_sf, fc2, _ = inputs(129) + fc2 = CountedWeight(fc2.shape, fc2.dtype) + fc2_sf = Tensor((1, 688128), "uint8") + run = self.ns["run_shared_fc12"] + for tokens in (129, 200): + run(Tensor((tokens, 7168)), fc1, fc1_sf, fc2, fc2_sf, swiglu_limit=10.0, sm_count=72) + # The second call of the same weight set skips the weight checks. + self.assertEqual(len(checks), 1) + layouts = self.ns["_WEIGHT_LAYOUTS"] + self.assertEqual(len(layouts), 1) + + # New weight tensors are validated again; entries of freed weights are pruned. + fc2_new = CountedWeight(fc2.shape, fc2.dtype) + del fc2 + run(activation, fc1, fc1_sf, fc2_new, fc2_sf, swiglu_limit=10.0, sm_count=72) + self.assertEqual(len(checks), 2) + self.assertEqual(len(layouts), 1) + + def test_tuning_key_text_is_unchanged(self): + # Persisted AutoTuner caches key the runner by str(unique_id()); the prebuilt + # string must stay the text of the identity tuple. + r = self.runner() + self.assertEqual( + r.unique_id(), + str( + ( + self.ns["_CATALOG_VERSION"], + 0, + 10, + 7, + 80, + 72, + 7168, + 3072, + 8192, + 10.0, + "K128_input_K32_intermediate_R128c4_gateup16", + ("fixed_weight_layout",), + ) + ), + ) + def test_terminal_release_drops_all_device_runner_references(self): first = self.runner(stream=17) second = self.runner(stream=18) From 37a1bc6358a52094a51666373dcb681136d35757 Mon Sep 17 00:00:00 2001 From: jintaop Date: Sun, 25 Oct 2026 20:49:08 -0700 Subject: [PATCH 40/59] [None][fix] Raise the HALO-Q serving spin bound Default the HALO-Q scheduler's serving spin bound to 300 s, the same order as NVLinkOneSided's native serving timeout, so a transient stall on one rank does not trap every EP rank. TRTLLM_MOE_COMM_TIMEOUT_SEC still selects a shorter bound. Signed-off-by: jintaop --- .../megamoe_scheduler/cuda_scheduler/runtime.py | 13 ++++++++----- .../_torch/moe/test_dynamic_eplb_contracts.py | 2 +- .../_torch/moe/test_per_layer_dynamic_eplb.py | 6 +++--- 3 files changed, 12 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py index eddea7afada0..cd76eb9e6ef9 100644 --- a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py @@ -16,13 +16,16 @@ MAX_EXPERTS = 384 MAX_CTAS = 128 MAX_BROADCASTS = MAX_EXPERTS -# Default bound, in SM clock cycles, of the scheduler's grid barrier and -# cross-rank waits; a timeout records a status flag. Each launch takes its bound -# as an argument, so callers can allow more peer skew while ranks warm up. -SPIN_CYCLES = 8_000_000_000 # Nominal clock64 rate that converts a timeout in seconds to SM cycles, as the -# MoE all-to-all kernels assume. SPIN_CYCLES is about 4 s at this rate. +# MoE all-to-all kernels assume. ASSUMED_CLOCK_HZ = 2_000_000_000 +# Default bound, in SM clock cycles, of the scheduler's grid barrier and +# cross-rank waits; a timeout records a status flag. Each launch takes its bound +# as an argument, so callers can allow more peer skew while ranks warm up. The +# default matches NVLinkOneSided's 300 s serving timeout, so a transient stall +# on one rank does not trap every EP rank; TRTLLM_MOE_COMM_TIMEOUT_SEC can set +# a shorter bound. +SPIN_CYCLES = 300 * ASSUMED_CLOCK_HZ def spin_cycles_for_timeout(seconds: int | None) -> int: diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index 46dc62d845f9..74584df8e092 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -28,7 +28,7 @@ _SCHEDULER_RUNTIME = _TORCH_ROOT / "moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py" # The scheduler's native bound and the bound of the default 1800 s MoE # communication warmup timeout at the nominal 2 GHz clock. -_SERVING_SPIN_CYCLES = 8_000_000_000 +_SERVING_SPIN_CYCLES = 300 * 2_000_000_000 _WARMUP_SPIN_CYCLES = 1800 * 2_000_000_000 diff --git a/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py b/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py index 1e9836aa00be..672d4dfbcc29 100644 --- a/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py +++ b/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py @@ -479,18 +479,18 @@ def test_scheduler_spin_bound_follows_the_warmup_guard_until_shutdown(monkeypatc bound = manager.get_scheduler_spin_bound() assert manager.get_scheduler_spin_bound() is bound - assert bound.cycles() == 8_000_000_000 + assert bound.cycles() == 300 * 2_000_000_000 tuner.is_tuning_mode = True assert bound.cycles() == 1800 * 2_000_000_000 tuner.is_tuning_mode = False guard.set_warmup(True) assert bound.cycles() == 1800 * 2_000_000_000 guard.set_warmup(False) - assert bound.cycles() == 8_000_000_000 + assert bound.cycles() == 300 * 2_000_000_000 assert manager.shutdown() guard.set_warmup(True) # A released bound no longer follows the guard. - assert bound.cycles() == 8_000_000_000 + assert bound.cycles() == 300 * 2_000_000_000 def test_per_layer_dynamic_manager_shares_the_base_lifecycle_only(): From d04156f3af7ca5a7e3ff207a52f6c84fe711c435 Mon Sep 17 00:00:00 2001 From: jintaop Date: Sun, 25 Oct 2026 23:54:43 -0700 Subject: [PATCH 41/59] [None][fix] Limit collective executor shutdown to per-layer dynamic EPLB Only the per-layer dynamic EPLB fabric arena needs an all-rank release. Gate the all-rank allgathers, the raise-and-retry handshake and the RPC two-phase commit on moe_config.load_balancer.mode == "per_layer_dynamic"; every other configuration keeps the best-effort, rank-local teardown of the MPI, RPC and Ray workers, plus a rank-local release of shared-FC12 runners. With per-layer dynamic EPLB: - a dead Ray actor falls through to forced teardown, since it can never join the collective release; - the RPC cleanup and commit calls are bounded at 300 s, so a dead rank zero cannot hang an atexit shutdown; - a finalizer no longer retries a failed collective shutdown; only an explicit call does. Also write the HALO-Q serving spin bound as a literal, which the scheduler contract tests parse with ast.literal_eval. Signed-off-by: jintaop --- .../cuda_scheduler/runtime.py | 8 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 25 +++ tensorrt_llm/executor/base_worker.py | 28 ++- tensorrt_llm/executor/ray/executor.py | 94 +++++---- tensorrt_llm/executor/ray/gpu_worker.py | 50 ++++- tensorrt_llm/executor/rpc_proxy.py | 75 +++++++- tensorrt_llm/executor/rpc_worker.py | 12 +- tensorrt_llm/executor/worker.py | 73 ++++++- tensorrt_llm/llmapi/_load_balance_env.py | 11 ++ .../_torch/moe/test_dynamic_eplb_contracts.py | 14 +- .../executor/test_shutdown_retry_contracts.py | 181 ++++++++++++++++-- 11 files changed, 505 insertions(+), 66 deletions(-) diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py index cd76eb9e6ef9..671fece56427 100644 --- a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py @@ -22,10 +22,10 @@ # Default bound, in SM clock cycles, of the scheduler's grid barrier and # cross-rank waits; a timeout records a status flag. Each launch takes its bound # as an argument, so callers can allow more peer skew while ranks warm up. The -# default matches NVLinkOneSided's 300 s serving timeout, so a transient stall -# on one rank does not trap every EP rank; TRTLLM_MOE_COMM_TIMEOUT_SEC can set -# a shorter bound. -SPIN_CYCLES = 300 * ASSUMED_CLOCK_HZ +# default, 300 s at ASSUMED_CLOCK_HZ, matches NVLinkOneSided's serving timeout, +# so a transient stall on one rank does not trap every EP rank; +# TRTLLM_MOE_COMM_TIMEOUT_SEC can set a shorter bound. +SPIN_CYCLES = 600_000_000_000 def spin_cycles_for_timeout(seconds: int | None) -> int: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 356652978015..bffe50156091 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2363,6 +2363,31 @@ def _release_execution_resources(self, self.dwdp_manager = None return True + def local_cleanup(self) -> None: + """Rank-local teardown for executors without collective EP resources. + + Drops the terminal model snapshot so the engines release their model + state through their own finalizers, as before collective cleanup. + """ + try: + self.release_rank_local_caches() + finally: + self._terminal_model_engines = () + + def release_rank_local_caches(self) -> None: + """Release this executor stream's shared-FC12 runners, if any were built.""" + shared_fc12 = sys.modules.get( + "tensorrt_llm._torch.moe.kernels.megamoe_shared_fc12") + if shared_fc12 is None: + return + + device = self.execution_stream.device + device_index = device.index if hasattr(device, "index") else int(device) + shared_fc12.release_shared_fc12_cache( + device_index, + stream_handle=int(self.execution_stream.cuda_stream), + ) + def can_enqueue_requests(self) -> bool: """ Indicates if the current process is allowed to enqueue requests diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 6cb34992c25c..36a46780ffa8 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -37,7 +37,8 @@ nvtx_range_debug) from ..bindings import executor as tllm from ..bindings.BuildInfo import check_cuda_architecture_supported -from ..llmapi._load_balance_env import configure_moe_launch_queues +from ..llmapi._load_balance_env import (configure_moe_launch_queues, + uses_collective_shutdown) from ..llmapi.llm_args import BaseLlmArgs, ExecutorMemoryType from ..llmapi.tokenizer import TokenizerBase from ..llmapi.tracer import global_tracer @@ -1077,6 +1078,9 @@ def wakeup(self, wakeup_tags: list[str]) -> None: def shutdown(self): if self.doing_shutdown: return + if not uses_collective_shutdown(getattr(self, "llm_args", None)): + self._shutdown_rank_local() + return self.doing_shutdown = True try: @@ -1119,8 +1123,26 @@ def shutdown(self): except BaseException: # Keep the engine reachable so a caller can retry a failed # collective terminal cleanup in this long-lived worker process. + # Only an explicit call retries it; see __del__. self.doing_shutdown = False + self._collective_shutdown_failed = True raise + self._collective_shutdown_failed = False + + def _shutdown_rank_local(self): + """Best-effort teardown for configurations without collective EP resources.""" + self.doing_shutdown = True + + if self.engine is not None: + can_shutdown = getattr(self.engine, "can_shutdown", None) + if can_shutdown is None: + can_shutdown = self.engine.can_enqueue_requests + if can_shutdown(): + self.engine.shutdown() + local_cleanup = getattr(self.engine, "local_cleanup", None) + if local_cleanup is not None: + local_cleanup() + self.engine = None def get_disaggregated_params(self) -> dict: if self.engine is None or self.engine.kv_cache_transceiver is None: @@ -1297,6 +1319,10 @@ def __exit__(self, exc_type, exc_value, traceback) -> bool: return exc_type is None or exc_type == self.WorkerExit def __del__(self): + # Finalizers run at a different time on each rank, so they must not + # retry a failed collective shutdown. + if getattr(self, "_collective_shutdown_failed", False): + return self.shutdown() diff --git a/tensorrt_llm/executor/ray/executor.py b/tensorrt_llm/executor/ray/executor.py index d3dd2315422b..844f978d14b0 100644 --- a/tensorrt_llm/executor/ray/executor.py +++ b/tensorrt_llm/executor/ray/executor.py @@ -37,6 +37,7 @@ from tensorrt_llm.executor.ray.utils import unwrap_ray_errors from tensorrt_llm.logger import logger +from ...llmapi._load_balance_env import uses_collective_shutdown from ...llmapi.utils import logger_debug from ..executor import GenerationExecutor from ..postproc_worker import PostprocWorkerConfig @@ -364,41 +365,22 @@ def shutdown(self): # Then, shutdown the workers if hasattr(self, 'workers') and self.workers is not None: - shutdown_refs = getattr(self, "_pending_shutdown_refs", None) - if shutdown_refs is None: - shutdown_refs = [ - worker.shutdown.remote() for worker in self.workers - ] - self._pending_shutdown_refs = shutdown_refs - - # Converge every rank before deciding whether actor teardown is safe. - # Keep timed-out refs for the next explicit shutdown() call: issuing a - # second collective shutdown while the first is still running can - # deadlock the rank set. - shutdown_errors = [] - timed_out = False - deadline = time.monotonic() + 30.0 - for rank, shutdown_ref in enumerate(shutdown_refs): + if uses_collective_shutdown( + getattr(self, "worker_kwargs", {}).get("llm_args")): + self._shutdown_workers_collectively() + else: try: - ray.get(shutdown_ref, - timeout=max(0.0, deadline - time.monotonic())) - except ray.exceptions.GetTimeoutError as error: - timed_out = True - shutdown_errors.append((rank, error)) - except Exception as error: - shutdown_errors.append((rank, error)) - - if not timed_out: - self._pending_shutdown_refs = None - if shutdown_errors: - for rank, error in shutdown_errors: + shutdown_refs = [ + worker.shutdown.remote() for worker in self.workers + ] + # Add timeout to prevent indefinite hanging + ray.get(shutdown_refs, timeout=30.0) + except ray.exceptions.GetTimeoutError: logger.warning( - f"Ray worker {rank} shutdown failed: {error}") - # Preserve the actors and placement group. RayGPUWorker resets - # its own shutdown guard on failure, so a later explicit call can - # retry after all pending refs have converged. - self._shutdown_event.clear() - raise shutdown_errors[0][1] + "Timeout waiting for workers to shutdown after 30 seconds" + ) + except Exception as e: + logger.warning(f"Error shutting down: {e}") # The engines are already stopped by the shutdown RPC above, so # kill the actor processes explicitly instead of relying on @@ -436,6 +418,52 @@ def shutdown(self): logger.debug("Shutting down Ray cluster") ray.shutdown() + def _shutdown_workers_collectively(self): + """Shut down workers that must release EP resources collectively. + + Raises if a live worker failed, keeping the actors and placement group + for an explicit retry. A dead actor can never join the collective + release, so it falls through to forced teardown instead. + """ + shutdown_refs = getattr(self, "_pending_shutdown_refs", None) + if shutdown_refs is None: + shutdown_refs = [ + worker.shutdown.remote() for worker in self.workers + ] + self._pending_shutdown_refs = shutdown_refs + + # Converge every rank before deciding whether actor teardown is safe. + # Keep timed-out refs for the next explicit shutdown() call: issuing a + # second collective shutdown while the first is still running can + # deadlock the rank set. + shutdown_errors = [] + timed_out = False + actor_died = False + deadline = time.monotonic() + 30.0 + for rank, shutdown_ref in enumerate(shutdown_refs): + try: + ray.get(shutdown_ref, + timeout=max(0.0, deadline - time.monotonic())) + except ray.exceptions.GetTimeoutError as error: + timed_out = True + shutdown_errors.append((rank, error)) + except ray.exceptions.RayActorError as error: + actor_died = True + shutdown_errors.append((rank, error)) + except Exception as error: + shutdown_errors.append((rank, error)) + + if not timed_out or actor_died: + self._pending_shutdown_refs = None + for rank, error in shutdown_errors: + logger.warning(f"Ray worker {rank} shutdown failed: {error}") + if shutdown_errors and not actor_died: + # Preserve the actors and placement group. RayGPUWorker resets its + # own shutdown guard on failure, so a later explicit call can retry + # after all pending refs have converged. + self._shutdown_event.clear() + raise shutdown_errors[0][1] + def _wait_for_cluster_resource_release(self, timeout: float = 30.0) -> None: """Block until Ray returns the workers' resources to the cluster. diff --git a/tensorrt_llm/executor/ray/gpu_worker.py b/tensorrt_llm/executor/ray/gpu_worker.py index f2033c1d939e..2b634206d756 100644 --- a/tensorrt_llm/executor/ray/gpu_worker.py +++ b/tensorrt_llm/executor/ray/gpu_worker.py @@ -30,7 +30,8 @@ control_action_decorator) from ... import TorchLlmArgs -from ...llmapi._load_balance_env import configure_moe_launch_queues +from ...llmapi._load_balance_env import (configure_moe_launch_queues, + uses_collective_shutdown) from ...llmapi.llm_args import BaseLlmArgs, ExecutorMemoryType from ...llmapi.tokenizer import TokenizerBase from ...llmapi.utils import configure_cpu_affinity @@ -388,6 +389,9 @@ def shutdown(self): if self.doing_shutdown: return + if not uses_collective_shutdown(getattr(self, "llm_args", None)): + self._shutdown_rank_local() + return self.doing_shutdown = True shutdown_succeeded = False @@ -467,11 +471,55 @@ def shutdown(self): shutdown_succeeded = True logger.debug(f"Worker {self.rank} shutdown done.") finally: + # A failed collective shutdown can be retried only explicitly. + self._collective_shutdown_failed = not shutdown_succeeded if not shutdown_succeeded: # Preserve the engine and allow an explicit retry to finish any # model-owned collective teardown that could not complete. self.doing_shutdown = False + def _shutdown_rank_local(self): + """Best-effort teardown for configurations without collective EP resources.""" + self.doing_shutdown = True + + logger.debug(f'Worker {self.rank} shutting down...') + + if hasattr(self, 'shutdown_event'): + self.shutdown_event.set() + + if getattr(self, '_postproc_pool', None) is not None: + self.shutdown_postproc_workers() + + if hasattr(self, 'rpc_server') and self.rpc_server is not None: + logger.info(f"[Rank {self.global_rank}] Shutting down RPC server") + try: + self.rpc_server.shutdown() + except Exception as e: + # Suppress errors during RPC server shutdown + # These can occur if the server is already closed or during cleanup + logger.debug( + f"[Rank {self.global_rank}] Suppressed error during RPC server shutdown: {e}" + ) + self.rpc_server = None + + if self.engine is not None: + self.engine.shutdown() + local_cleanup = getattr(self.engine, "local_cleanup", None) + if local_cleanup is not None: + local_cleanup() + self.engine = None + + if (self.llm_args.backend == "pytorch" + and hasattr(self, "checkpoint_loader") + and self.checkpoint_loader is not None): + self.checkpoint_loader.cleanup() + self.checkpoint_loader = None + + # Check if there are any errors from the threads before shutdown. + self._handle_background_error() + + logger.debug(f"Worker {self.rank} shutdown done.") + def _configure_device(self): configure_cpu_affinity(self.device_id) diff --git a/tensorrt_llm/executor/rpc_proxy.py b/tensorrt_llm/executor/rpc_proxy.py index 60dd334b26c1..8606d38c075b 100644 --- a/tensorrt_llm/executor/rpc_proxy.py +++ b/tensorrt_llm/executor/rpc_proxy.py @@ -16,6 +16,7 @@ import threading from typing import List, Optional, Union +from ..llmapi._load_balance_env import uses_collective_shutdown from ..llmapi.mpi_session import (MpiPoolSession, MpiSession, validate_session_world_size) from ..llmapi.utils import logger_debug, print_colored @@ -32,6 +33,9 @@ class GenerationExecutorRpcProxy(RpcExecutorMixin, GenerationExecutor): + # Bound on each acknowledged RPC of the collective shutdown handshake. + _COLLECTIVE_SHUTDOWN_TIMEOUT_S = 300.0 + def __init__( self, worker_kwargs: dict, @@ -295,13 +299,23 @@ def get_data_transceiver_state(self) -> bytes: def setup_engine_remote(self): return self.rpc_client.setup_engine().remote(need_response=True) + def _uses_collective_shutdown(self) -> bool: + return uses_collective_shutdown( + getattr(self, "worker_kwargs", {}).get("llm_args")) + def shutdown_remote(self): logger_debug("Shutting down rpc remote", color="yellow") + if not self._uses_collective_shutdown(): + self.rpc_client.shutdown().remote(need_response=False) + return # Cleanup must be acknowledged before main_task is allowed to close the # server. The second RPC commits that acknowledged all-rank result; its - # response is sent before RpcWorker wakes main_task. - self.rpc_client.shutdown().remote(need_response=True) - self.rpc_client.commit_shutdown().remote(need_response=True) + # response is sent before RpcWorker wakes main_task. The timeout keeps + # a dead rank zero from hanging the caller, e.g. an atexit shutdown. + timeout = self._COLLECTIVE_SHUTDOWN_TIMEOUT_S + self.rpc_client.shutdown().remote(timeout=timeout, need_response=True) + self.rpc_client.commit_shutdown().remote(timeout=timeout, + need_response=True) def abort_request(self, request_id: int) -> None: return self.rpc_client.abort_request(request_id).remote() @@ -323,6 +337,9 @@ def stop_profile(self) -> None: def shutdown(self): if self._shutdown_event.is_set(): return + if not self._uses_collective_shutdown(): + self._shutdown_rank_local() + return self._shutdown_event.set() logger_debug("Shutting down GenerationExecutorRpcProxy", color="yellow") @@ -379,6 +396,58 @@ def shutdown(self): if shutdown_errors: raise shutdown_errors[0] + def _shutdown_rank_local(self): + """Best-effort teardown for configurations without collective EP resources.""" + self._shutdown_event.set() + logger_debug("Shutting down GenerationExecutorRpcProxy", color="yellow") + + # 1. shutdown the rpc server (PyExecutor Rank 0 + RPC server) + self.shutdown_remote() + + # 2. stop the main loop, so that no new rpc requests + if self.main_loop and self.main_loop_task_obj: + logger_debug("Cancelling main loop task.", color="yellow") + # The cancel() is thread-safe + try: + self.main_loop.call_soon_threadsafe( + self.main_loop_task_obj.cancel) + except Exception as e: + logger_debug(f"Error cancelling main loop task: {e}", + color="yellow") + + # Only join if we're not calling from the main_loop_thread itself + # (e.g., during garbage collection in that thread) + if self.main_loop_thread and threading.current_thread( + ) != self.main_loop_thread: + self.main_loop_thread.join(timeout=2.0) + if self.main_loop_thread.is_alive(): + logger.warning("Main loop thread did not exit gracefully") + + # 3. shutdown the mpi session, this should wait until all the PyExecutor + # processes are shutdown + if self.mpi_session is not None: + if self._owns_mpi_session: + logger_debug("Shutting down mpi session", color="yellow") + self.mpi_session.shutdown() + else: + # Externally owned (shared) session: leave the pool alive, but + # wait for this executor's worker tasks to finish so the next + # LLM on the pool doesn't race with PyExecutor teardown. Block + # without a timeout, mirroring the owned path above + # (mpi_session.shutdown() also waits indefinitely); a timeout + # that expires would just reintroduce the race it prevents. + for future in getattr(self, "worker_futures", []): + try: + future.result() + except Exception as e: + logger.warning( + f"RPC worker task raised during shutdown on " + f"shared MPI session: {e}") + logger_debug("Mpi session shutdown", color="yellow") + self.mpi_session = None + + self.rpc_client.close() + def __enter__(self): return self diff --git a/tensorrt_llm/executor/rpc_worker.py b/tensorrt_llm/executor/rpc_worker.py index 28c638e4f576..6f25cf71d240 100644 --- a/tensorrt_llm/executor/rpc_worker.py +++ b/tensorrt_llm/executor/rpc_worker.py @@ -11,6 +11,7 @@ from tensorrt_llm.llmapi.utils import enable_llm_debug, logger_debug from .._utils import mpi_rank +from ..llmapi._load_balance_env import uses_collective_shutdown from ..llmapi.llm_args import BaseLlmArgs from ..llmapi.tokenizer import TokenizerBase from ..logger import set_level @@ -105,8 +106,8 @@ def shutdown(self): logger_debug(f"[worker] RpcWorker #{mpi_rank()} is shutting down", color="yellow") if not getattr(self, "_rpc_collective_shutdown", False): - super().shutdown() self.shutdown_event.set() + super().shutdown() return # Followers wait between attempts. Wake them only when rank zero is @@ -187,7 +188,9 @@ def main_task( hf_model_dir=hf_model_dir, tokenizer=tokenizer, ) - worker._rpc_collective_shutdown = True + # Only the per-layer dynamic EPLB fabric arena needs the all-rank + # two-phase shutdown; other configurations keep rank-local teardown. + worker._rpc_collective_shutdown = uses_collective_shutdown(llm_args) if mpi_rank() != 0: # The non-leader worker will setup the engine immediately. @@ -200,8 +203,9 @@ def main_task( # Rank zero publishes the shutdown request. Every follower waits # for its broadcast sentinel, then runs the same finalize path so # model-owned EP collectives remain rank-aligned. - worker.engine.wait_shutdown() - while True: + if worker._rpc_collective_shutdown: + worker.engine.wait_shutdown() + while worker._rpc_collective_shutdown: shutdown_error = None try: worker.shutdown() diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py index 459c8867f53c..68848e5bc473 100644 --- a/tensorrt_llm/executor/worker.py +++ b/tensorrt_llm/executor/worker.py @@ -28,7 +28,8 @@ from tensorrt_llm.logger import logger from .._utils import mpi_comm, mpi_rank, print_all_stacks -from ..llmapi._load_balance_env import configure_moe_launch_queues +from ..llmapi._load_balance_env import (configure_moe_launch_queues, + uses_collective_shutdown) from ..llmapi.llm_args import BaseLlmArgs from ..llmapi.mpi_session import set_mpi_session_cpp from ..llmapi.tokenizer import TokenizerBase @@ -154,6 +155,9 @@ def shutdown(self): if self.doing_shutdown: return + if not uses_collective_shutdown(getattr(self, "llm_args", None)): + self._shutdown_rank_local() + return self.doing_shutdown = True shutdown_succeeded = False @@ -276,9 +280,76 @@ def shutdown(self): # every rank has reached, so a retry cannot re-enter collectives. self._handle_background_error() finally: + # A failed collective shutdown can be retried only explicitly. + self._collective_shutdown_failed = not shutdown_succeeded if not shutdown_succeeded: self.doing_shutdown = False + def _shutdown_rank_local(self): + """Best-effort teardown for configurations without collective EP resources.""" + self.doing_shutdown = True + + logger_debug(f'Worker {mpi_rank()} shutdown...\n', "yellow") + + # The server references this worker; a long-lived worker process would + # otherwise keep one server, its threads and sockets per LLM it hosted. + if self.rpc_server is not None: + self.rpc_server.shutdown() + self.rpc_server = None + + if self.engine is not None: + if self.engine.can_enqueue_requests(): + if self.await_response_thread.is_alive(): + self.await_response_thread.stop() + self.await_response_thread.join() + + self.engine.shutdown() + local_cleanup = getattr(self.engine, "local_cleanup", None) + if local_cleanup is not None: + local_cleanup() + self.engine = None + + if (self.llm_args.backend == "pytorch" + and hasattr(self, "checkpoint_loader") + and self.checkpoint_loader is not None): + self.checkpoint_loader.cleanup() + self.checkpoint_loader = None + + # MegaMoE's NVLink symmetric-memory activation workspaces are + # rendezvoused over the EP group, so they must go before the + # destroy_process_group() below. Never let a failure here escape: + # doing_shutdown is already set, so a retry would no-op, and skipping + # the teardown below would leave NCCL communicators alive -- turning a + # MegaMoE-only failure into a worker teardown hang. + mega_moe = sys.modules.get(_MEGA_MOE_DEEPGEMM_MODULE) + if mega_moe is not None: + try: + mega_moe.release_symm_buffer_cache() + except Exception as e: + logger.error( + f"Failed to release MegaMoE symm buffers on shutdown: {e}") + + # Destroy torch distributed process groups so that NCCL communicators + # are torn down cleanly before MPI session shutdown and process exit. + # This is done here (not in PyExecutor.shutdown()) because the MPI + # worker owns the process group. In the Ray path the process group + # belongs to RayWorkerWrapper and must not be destroyed by the engine. + import torch.distributed + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + # Return this rank's GPU memory to the driver. Under an external MPI + # launch (mpirun/srun, e.g. CI), the worker process is long-lived and + # shared across successive LLM instances, so the previous model's + # reservation would otherwise carry into the next LLM in this process. + gc.collect() + torch.cuda.empty_cache() + + # Check if there are any errors from the threads before shutdown. + self._handle_background_error() + + logger_debug(f"Worker {mpi_rank()} shutdown done.\n", "yellow") + def block_subordinates(self): if self.rank != 0: from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor diff --git a/tensorrt_llm/llmapi/_load_balance_env.py b/tensorrt_llm/llmapi/_load_balance_env.py index 5fd77b9d4b62..cfb12e65eeb4 100644 --- a/tensorrt_llm/llmapi/_load_balance_env.py +++ b/tensorrt_llm/llmapi/_load_balance_env.py @@ -47,3 +47,14 @@ def configure_moe_launch_queues( overrides[name] = "4x" os.environ[name] = "4x" return overrides + + +def uses_collective_shutdown(llm_args: Any) -> bool: + """Return whether executor teardown must release EP resources collectively. + + Only the per-layer dynamic EPLB fabric arena needs an all-rank release. + Every other configuration keeps the best-effort, rank-local teardown. + """ + moe_config = getattr(llm_args, "moe_config", None) + load_balancer = getattr(moe_config, "load_balancer", None) + return getattr(load_balancer, "mode", None) == "per_layer_dynamic" diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index 74584df8e092..d71c20718973 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -23,6 +23,9 @@ pytestmark = pytest.mark.cpu_only _ROOT = Path(__file__).resolve().parents[4] +# Worker shutdown contracts below cover the collective path, which only +# per-layer dynamic EPLB takes. +_PER_LAYER_DYNAMIC_SHUTDOWN = {"uses_collective_shutdown": lambda llm_args: True} _TORCH_ROOT = _ROOT / "tensorrt_llm/_torch" _QUEUE = "CUDA_SCALE_LAUNCH_QUEUES" _SCHEDULER_RUNTIME = _TORCH_ROOT / "moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py" @@ -756,10 +759,11 @@ def test_executor_terminal_cleanup_is_three_phase_and_stream_scoped(monkeypatch) executor_cls = _extract_class( _TORCH_ROOT / "pyexecutor/py_executor.py", "PyExecutor", - {"terminal_cleanup"}, + {"terminal_cleanup", "release_rank_local_caches"}, { "__name__": "tensorrt_llm._torch.pyexecutor.py_executor", "__package__": "tensorrt_llm._torch.pyexecutor", + "sys": sys, }, ) @@ -899,7 +903,7 @@ def __exit__(self, *args): _ROOT / "tensorrt_llm/executor/base_worker.py", "BaseWorker", {"shutdown"}, - {}, + _PER_LAYER_DYNAMIC_SHUTDOWN, ) engine = Engine() @@ -2485,7 +2489,7 @@ def test_base_worker_finalizes_rank_follower_pyexecutor(): _ROOT / "tensorrt_llm/executor/base_worker.py", "BaseWorker", {"shutdown"}, - {}, + _PER_LAYER_DYNAMIC_SHUTDOWN, ) events = [] engine = SimpleNamespace( @@ -2565,6 +2569,7 @@ def test_mpi_worker_terminal_cleanup_failure_is_retryable(monkeypatch): "GenerationExecutorWorker", {"shutdown"}, { + **_PER_LAYER_DYNAMIC_SHUTDOWN, "logger_debug": lambda *args, **kwargs: None, "mpi_rank": lambda: 0, "mpi_comm": lambda: SimpleNamespace(allgather=lambda value: [value]), @@ -2622,6 +2627,7 @@ def test_ray_worker_postproc_failure_still_runs_unsafe_terminal_cleanup(): "RayGPUWorker", {"shutdown"}, { + **_PER_LAYER_DYNAMIC_SHUTDOWN, "logger": SimpleNamespace( debug=lambda *args, **kwargs: None, info=lambda *args, **kwargs: None ), @@ -2679,6 +2685,7 @@ def test_ray_worker_shutdown_preserves_first_error_and_attempts_all_cleanup(): "RayGPUWorker", {"shutdown"}, { + **_PER_LAYER_DYNAMIC_SHUTDOWN, "logger": SimpleNamespace( debug=lambda *args, **kwargs: None, info=lambda *args, **kwargs: None ), @@ -2728,6 +2735,7 @@ def test_ray_worker_partial_init_without_llm_args_still_checks_background(): "RayGPUWorker", {"shutdown"}, { + **_PER_LAYER_DYNAMIC_SHUTDOWN, "logger": SimpleNamespace( debug=lambda *args, **kwargs: None, info=lambda *args, **kwargs: None ), diff --git a/tests/unittest/executor/test_shutdown_retry_contracts.py b/tests/unittest/executor/test_shutdown_retry_contracts.py index f5485fd8c33d..f9bb440e5235 100644 --- a/tests/unittest/executor/test_shutdown_retry_contracts.py +++ b/tests/unittest/executor/test_shutdown_retry_contracts.py @@ -50,10 +50,31 @@ def _extract_class(source, class_name, methods, namespace, bases=()): return namespace[class_name] +_uses_collective_shutdown = _extract_function( + _ROOT / "tensorrt_llm/llmapi/_load_balance_env.py", "uses_collective_shutdown", {"Any": object} +) + + +def _llm_args(per_layer_dynamic=True, **fields): + load_balancer = SimpleNamespace(mode="per_layer_dynamic") if per_layer_dynamic else None + return SimpleNamespace(moe_config=SimpleNamespace(load_balancer=load_balancer), **fields) + + +def test_collective_shutdown_is_limited_to_per_layer_dynamic_eplb(): + assert _uses_collective_shutdown(_llm_args()) + assert not _uses_collective_shutdown(_llm_args(per_layer_dynamic=False)) + assert not _uses_collective_shutdown(SimpleNamespace(backend="pytorch")) + assert not _uses_collective_shutdown(None) + + class _Timeout(Exception): pass +class _ActorDied(Exception): + pass + + class _Ref: def __init__(self, outcomes): self.outcomes = list(outcomes) @@ -83,7 +104,7 @@ def __init__(self, refs): class _FakeRay: def __init__(self): - self.exceptions = SimpleNamespace(GetTimeoutError=_Timeout) + self.exceptions = SimpleNamespace(GetTimeoutError=_Timeout, RayActorError=_ActorDied) self.util = SimpleNamespace(remove_placement_group=lambda group: None) self.killed = [] @@ -103,13 +124,14 @@ def shutdown(): return None -def _ray_executor(ray): +def _ray_executor(ray, per_layer_dynamic=True): cls = _extract_class( _ROOT / "tensorrt_llm/executor/ray/executor.py", "RayExecutor", - {"shutdown"}, + {"shutdown", "_shutdown_workers_collectively"}, { "ray": ray, + "uses_collective_shutdown": _uses_collective_shutdown, "time": __import__("time"), "logger": SimpleNamespace( warning=lambda *args, **kwargs: None, debug=lambda *args, **kwargs: None @@ -118,6 +140,7 @@ def _ray_executor(ray): }, ) executor = cls() + executor.worker_kwargs = {"llm_args": _llm_args(per_layer_dynamic)} executor._shutdown_event = threading.Event() executor.main_loop = None executor.main_loop_task_obj = None @@ -164,6 +187,32 @@ def test_ray_shutdown_timeout_reuses_pending_ref(): assert executor.workers is None +def test_ray_dead_actor_falls_through_to_forced_teardown(): + ray = _FakeRay() + dead = _RayWorker([_Ref([_ActorDied("actor died")])]) + live = _RayWorker([_Ref([None])]) + executor = _ray_executor(ray) + executor.workers = [dead, live] + + executor.shutdown() + + assert ray.killed == [dead, live] + assert executor.workers is None + assert executor._pending_shutdown_refs is None + + +def test_ray_shutdown_without_per_layer_eplb_tears_down_best_effort(): + ray = _FakeRay() + worker = _RayWorker([_Ref([RuntimeError("cleanup failed")])]) + executor = _ray_executor(ray, per_layer_dynamic=False) + executor.workers = [worker] + + executor.shutdown() + + assert ray.killed == [worker] + assert executor.workers is None + + class _BaseWorker: def shutdown(self): outcome = self.base_outcomes.pop(0) @@ -247,8 +296,8 @@ def __init__(self, name, calls, error=None): self.calls = calls self.error = error - def remote(self, *, need_response): - self.calls.append((self.name, need_response)) + def remote(self, *, need_response, timeout=None): + self.calls.append((self.name, need_response, timeout)) if self.error is not None: raise self.error @@ -265,35 +314,50 @@ def commit_shutdown(self): return _RemoteCall("commit", self.calls) def close(self): - self.calls.append(("close", True)) + self.calls.append(("close", True, None)) -def _rpc_proxy(): - return _extract_class( +def _rpc_proxy(per_layer_dynamic=True): + cls = _extract_class( _ROOT / "tensorrt_llm/executor/rpc_proxy.py", "GenerationExecutorRpcProxy", - {"shutdown_remote", "shutdown"}, + {"_uses_collective_shutdown", "shutdown_remote", "shutdown", "_shutdown_rank_local"}, { + "uses_collective_shutdown": _uses_collective_shutdown, "logger_debug": lambda *args, **kwargs: None, "logger": SimpleNamespace(warning=lambda *args, **kwargs: None), "threading": threading, }, ) + cls._COLLECTIVE_SHUTDOWN_TIMEOUT_S = 300.0 + proxy = cls() + proxy.worker_kwargs = {"llm_args": _llm_args(per_layer_dynamic)} + return proxy def test_rpc_proxy_requires_cleanup_and_commit_responses(): calls = [] - proxy = _rpc_proxy()() + proxy = _rpc_proxy() + proxy.rpc_client = _RpcClient(calls) + + proxy.shutdown_remote() + + assert calls == [("cleanup", True, 300.0), ("commit", True, 300.0)] + + +def test_rpc_proxy_without_per_layer_eplb_sends_one_unacknowledged_shutdown(): + calls = [] + proxy = _rpc_proxy(per_layer_dynamic=False) proxy.rpc_client = _RpcClient(calls) proxy.shutdown_remote() - assert calls == [("cleanup", True), ("commit", True)] + assert calls == [("cleanup", False, None)] def test_rpc_proxy_cleanup_failure_keeps_retry_gate_and_resources_live(): calls = [] - proxy = _rpc_proxy()() + proxy = _rpc_proxy() proxy._shutdown_event = threading.Event() proxy.rpc_client = _RpcClient(calls, RuntimeError("remote cleanup failed")) proxy.main_loop = None @@ -306,7 +370,7 @@ def test_rpc_proxy_cleanup_failure_keeps_retry_gate_and_resources_live(): assert not proxy._shutdown_event.is_set() assert proxy.mpi_session is not None - assert calls == [("cleanup", True)] + assert calls == [("cleanup", True, 300.0)] def test_rpc_server_completes_commit_only_after_response_send(): @@ -405,6 +469,7 @@ def release_symm_buffer_cache(): "GenerationExecutorWorker", {"shutdown"}, { + "uses_collective_shutdown": _uses_collective_shutdown, "logger_debug": lambda *args, **kwargs: None, "mpi_rank": lambda: 0, "mpi_comm": lambda: SimpleNamespace(allgather=lambda value: [value]), @@ -423,7 +488,7 @@ def release_symm_buffer_cache(): worker.doing_shutdown = False worker.rpc_server = None worker.engine = engine - worker.llm_args = SimpleNamespace(backend="pytorch") + worker.llm_args = _llm_args(backend="pytorch") worker._handle_background_error = lambda: events.append("background") with pytest.raises(RuntimeError, match="symm release failed"): @@ -486,6 +551,7 @@ def allgather(self, value): "GenerationExecutorWorker", {"shutdown"}, { + "uses_collective_shutdown": _uses_collective_shutdown, "logger_debug": lambda *args, **kwargs: None, "mpi_rank": lambda: 0, "mpi_comm": lambda: Comm(), @@ -504,7 +570,7 @@ def allgather(self, value): worker.doing_shutdown = False worker.rpc_server = None worker.engine = engine - worker.llm_args = SimpleNamespace(backend="pytorch") + worker.llm_args = _llm_args(backend="pytorch") worker._handle_background_error = lambda: events.append("background") with pytest.raises(RuntimeError, match="rank 1.*peer cleanup failed"): @@ -555,6 +621,7 @@ def test_mpi_worker_background_error_keeps_shutdown_committed(monkeypatch): "GenerationExecutorWorker", {"shutdown"}, { + "uses_collective_shutdown": _uses_collective_shutdown, "logger_debug": lambda *args, **kwargs: None, "mpi_rank": lambda: 0, "mpi_comm": lambda: SimpleNamespace(allgather=lambda value: [value]), @@ -572,7 +639,7 @@ def test_mpi_worker_background_error_keeps_shutdown_committed(monkeypatch): shutdown=lambda: None, terminal_cleanup=lambda *, local_safe: cleanups.append(local_safe), ) - worker.llm_args = SimpleNamespace(backend="pytorch") + worker.llm_args = _llm_args(backend="pytorch") def background_error(): raise RuntimeError("background") @@ -585,3 +652,85 @@ def background_error(): assert worker.doing_shutdown and worker.engine is None worker.shutdown() assert cleanups == [True] + + +def test_mpi_worker_without_per_layer_eplb_keeps_best_effort_teardown(monkeypatch): + events = [] + module_name = "test_best_effort_megamoe_symm_release" + + def release_symm_buffer_cache(): + events.append("release_symm") + raise RuntimeError("symm release failed") + + def no_collectives(): + raise AssertionError("rank-local teardown must not enter MPI collectives") + + mega_moe = types.ModuleType(module_name) + mega_moe.release_symm_buffer_cache = release_symm_buffer_cache + fake_distributed = types.ModuleType("torch.distributed") + fake_distributed.is_initialized = lambda: True + fake_distributed.destroy_process_group = lambda: events.append("destroy_pg") + fake_torch = types.ModuleType("torch") + fake_torch.distributed = fake_distributed + fake_torch.cuda = SimpleNamespace(empty_cache=lambda: events.append("empty_cache")) + monkeypatch.setitem(sys.modules, module_name, mega_moe) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + monkeypatch.setitem(sys.modules, "torch.distributed", fake_distributed) + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/worker.py", + "GenerationExecutorWorker", + {"shutdown", "_shutdown_rank_local"}, + { + "uses_collective_shutdown": _uses_collective_shutdown, + "logger_debug": lambda *args, **kwargs: None, + "logger": SimpleNamespace(error=lambda *args, **kwargs: events.append("logged")), + "mpi_rank": lambda: 0, + "mpi_comm": no_collectives, + "sys": sys, + "_MEGA_MOE_DEEPGEMM_MODULE": module_name, + "gc": SimpleNamespace(collect=lambda: events.append("gc")), + }, + ) + worker = worker_cls() + worker.doing_shutdown = False + worker.rpc_server = None + worker.engine = SimpleNamespace( + can_enqueue_requests=lambda: False, + shutdown=lambda: events.append("engine_shutdown"), + local_cleanup=lambda: events.append("local_cleanup"), + terminal_cleanup=lambda **kwargs: events.append("terminal_cleanup"), + ) + worker.llm_args = _llm_args(per_layer_dynamic=False, backend="pytorch") + worker._handle_background_error = lambda: events.append("background") + + worker.shutdown() + + assert events == [ + "engine_shutdown", + "local_cleanup", + "release_symm", + "logged", + "destroy_pg", + "gc", + "empty_cache", + "background", + ] + assert worker.doing_shutdown and worker.engine is None + + +def test_finalizer_does_not_retry_a_failed_collective_shutdown(): + calls = [] + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/base_worker.py", "BaseWorker", {"__del__"}, {} + ) + worker = worker_cls() + worker.shutdown = lambda: calls.append("shutdown") + + worker._collective_shutdown_failed = True + worker.__del__() + assert calls == [] + + worker._collective_shutdown_failed = False + worker.__del__() + assert calls == ["shutdown"] + worker.shutdown = lambda: None From febdbc83e713f052052f27d333d43073f90d0f4f Mon Sep 17 00:00:00 2001 From: jintaop Date: Mon, 26 Oct 2026 02:27:32 -0700 Subject: [PATCH 42/59] [None][fix] Keep standard EPLB's finalizer shutdown and check launch queues before overrides Standard EPLB is torn down rank-locally since collective shutdown is limited to per-layer dynamic EPLB, so restore main's MoeLoadBalancer.__del__ -> shutdown(), which unlinks the /dev/shm segments holding the host expert weights. The model-engine finalizer now warns about a skipped collective shutdown only for balancers that need one. The late-initialization check for CUDA_SCALE_LAUNCH_QUEUES=4x now reads the setting the process had before env_overrides were replayed: a CUDA context that already exists was created with that value, so a reused worker that initialized CUDA without 4x is rejected instead of running per-layer dynamic EPLB without the larger launch queues. Signed-off-by: jintaop --- .../_torch/moe/fused_moe/moe_load_balancer.py | 10 ++--- .../_torch/pyexecutor/model_engine.py | 25 +++++------ tensorrt_llm/executor/base_worker.py | 8 +++- tensorrt_llm/executor/ray/gpu_worker.py | 8 +++- tensorrt_llm/executor/worker.py | 8 +++- tensorrt_llm/llmapi/_load_balance_env.py | 16 +++++++- tensorrt_llm/llmapi/llm.py | 5 ++- .../_torch/moe/test_dynamic_eplb_contracts.py | 41 +++++++++++++++++-- 8 files changed, 88 insertions(+), 33 deletions(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py b/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py index a8f48b9b3062..4b3ac2c3b189 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py @@ -916,13 +916,13 @@ def __init__(self, self.next_layer_repeated_count = None + # Without per-layer dynamic EPLB, executor teardown is rank-local and this + # balancer is released by its own finalizer, which unlinks its shared memory. + shutdown_on_finalize = True + def __del__(self): - # MPI collectives are unsafe from rank-asymmetric Python finalization. - # Normal worker teardown calls shutdown() explicitly on every rank. if not getattr(self, "is_shutdown", True): - logger.warning( - "MoeLoadBalancer was finalized before explicit collective shutdown" - ) + self.shutdown() def is_static_routing(self): # if we don't update, then it is statistic routing. diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 2b90bcd9d917..7454ddca3fb0 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1105,14 +1105,11 @@ def invalidate_multimodal_encoder_cache(self) -> None: if encoder_cache is not None: encoder_cache.clear() - def prepare_cleanup(self, - *, - release_nccl_window_owners: bool = True) -> None: + def prepare_cleanup(self, *, release_nccl_window_owners: bool = True) -> None: """Release local graph captures before collective model teardown.""" if self._cleanup_done: return - self._release_cuda_graphs( - release_nccl_window_owners=release_nccl_window_owners) + self._release_cuda_graphs(release_nccl_window_owners=release_nccl_window_owners) def shutdown_moe_load_balancer(self, *, local_safe: bool) -> bool: """Run the manager's rank-aligned teardown preflight and collectives.""" @@ -1121,8 +1118,7 @@ def shutdown_moe_load_balancer(self, *, local_safe: bool) -> bool: moe_load_balancer = getattr(self, "moe_load_balancer", None) if moe_load_balancer is None: return bool(local_safe) - return moe_load_balancer.shutdown( - local_safe=bool(local_safe)) is not False + return moe_load_balancer.shutdown(local_safe=bool(local_safe)) is not False def finish_cleanup(self) -> None: """Release rank-local model resources after collective teardown.""" @@ -1163,10 +1159,14 @@ def cleanup(self, *, collective: bool = True) -> None: # Destruction timing is not rank-symmetric; only explicit shutdown # releases NCCL window owners. self.prepare_cleanup(release_nccl_window_owners=False) - if getattr(self, "moe_load_balancer", None) is not None: + moe_load_balancer = getattr(self, "moe_load_balancer", None) + if moe_load_balancer is not None and not getattr( + moe_load_balancer, "shutdown_on_finalize", False + ): logger.warning( "Skipping collective MoE load-balancer shutdown from Python finalization; " - "terminal executor shutdown must release it explicitly") + "terminal executor shutdown must release it explicitly" + ) self.finish_cleanup() return @@ -1176,14 +1176,11 @@ def cleanup(self, *, collective: bool = True) -> None: except BaseException as error: # Preserve rank alignment before surfacing it. prepare_error = error - manager_safe = self.shutdown_moe_load_balancer( - local_safe=prepare_error is None) + manager_safe = self.shutdown_moe_load_balancer(local_safe=prepare_error is None) if prepare_error is not None: raise prepare_error if not manager_safe: - raise RuntimeError( - "MoE load-balancer collective shutdown was not safe on every rank" - ) + raise RuntimeError("MoE load-balancer collective shutdown was not safe on every rank") self.finish_cleanup() def __del__(self) -> None: diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 36a46780ffa8..699fa7393ff2 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -37,7 +37,8 @@ nvtx_range_debug) from ..bindings import executor as tllm from ..bindings.BuildInfo import check_cuda_architecture_supported -from ..llmapi._load_balance_env import (configure_moe_launch_queues, +from ..llmapi._load_balance_env import (LAUNCH_QUEUES_ENV, + configure_moe_launch_queues, uses_collective_shutdown) from ..llmapi.llm_args import BaseLlmArgs, ExecutorMemoryType from ..llmapi.tokenizer import TokenizerBase @@ -102,11 +103,14 @@ def __init__( llm_args: Optional[BaseLlmArgs] = None, ) -> None: # Direct worker construction also bypasses BaseLLM/worker_main. + launch_queues = os.environ.get(LAUNCH_QUEUES_ENV) if llm_args is not None and llm_args.env_overrides: os.environ.update(llm_args.env_overrides) if llm_args is not None and llm_args.backend == "pytorch": queue_overrides = configure_moe_launch_queues( - getattr(llm_args, "moe_config", None), llm_args.env_overrides) + getattr(llm_args, "moe_config", None), + llm_args.env_overrides, + initial_setting=launch_queues) if queue_overrides is not llm_args.env_overrides: llm_args.env_overrides = queue_overrides diff --git a/tensorrt_llm/executor/ray/gpu_worker.py b/tensorrt_llm/executor/ray/gpu_worker.py index 2b634206d756..a2d8b34f7e65 100644 --- a/tensorrt_llm/executor/ray/gpu_worker.py +++ b/tensorrt_llm/executor/ray/gpu_worker.py @@ -30,7 +30,8 @@ control_action_decorator) from ... import TorchLlmArgs -from ...llmapi._load_balance_env import (configure_moe_launch_queues, +from ...llmapi._load_balance_env import (LAUNCH_QUEUES_ENV, + configure_moe_launch_queues, uses_collective_shutdown) from ...llmapi.llm_args import BaseLlmArgs, ExecutorMemoryType from ...llmapi.tokenizer import TokenizerBase @@ -59,11 +60,14 @@ class RayWorkerWrapper: def __init__(self, worker_cls, worker_kwargs, world_size, rank): llm_args = worker_kwargs.get("llm_args") + launch_queues = os.environ.get(LAUNCH_QUEUES_ENV) if llm_args is not None and llm_args.env_overrides: os.environ.update(llm_args.env_overrides) if llm_args is not None and llm_args.backend == "pytorch": queue_overrides = configure_moe_launch_queues( - getattr(llm_args, "moe_config", None), llm_args.env_overrides) + getattr(llm_args, "moe_config", None), + llm_args.env_overrides, + initial_setting=launch_queues) if queue_overrides is not llm_args.env_overrides: llm_args.env_overrides = queue_overrides diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py index 68848e5bc473..9c25aaf88f3f 100644 --- a/tensorrt_llm/executor/worker.py +++ b/tensorrt_llm/executor/worker.py @@ -28,7 +28,8 @@ from tensorrt_llm.logger import logger from .._utils import mpi_comm, mpi_rank, print_all_stacks -from ..llmapi._load_balance_env import (configure_moe_launch_queues, +from ..llmapi._load_balance_env import (LAUNCH_QUEUES_ENV, + configure_moe_launch_queues, uses_collective_shutdown) from ..llmapi.llm_args import BaseLlmArgs from ..llmapi.mpi_session import set_mpi_session_cpp @@ -376,6 +377,7 @@ def worker_main( hmac_key: bytes = b"", ) -> None: + launch_queues = os.environ.get(LAUNCH_QUEUES_ENV) if llm_args is not None and llm_args.env_overrides: # MPI may have cached the parent's environment before this worker was # spawned, so replay overrides before any rebalance state is inspected. @@ -384,7 +386,9 @@ def worker_main( # Configure per-layer dynamic EPLB before MPI or CUDA state is created. if llm_args is not None and llm_args.backend == "pytorch": queue_overrides = configure_moe_launch_queues( - getattr(llm_args, "moe_config", None), llm_args.env_overrides) + getattr(llm_args, "moe_config", None), + llm_args.env_overrides, + initial_setting=launch_queues) if queue_overrides is not llm_args.env_overrides: llm_args.env_overrides = queue_overrides diff --git a/tensorrt_llm/llmapi/_load_balance_env.py b/tensorrt_llm/llmapi/_load_balance_env.py index cfb12e65eeb4..8951b6b7cac8 100644 --- a/tensorrt_llm/llmapi/_load_balance_env.py +++ b/tensorrt_llm/llmapi/_load_balance_env.py @@ -15,10 +15,15 @@ if TYPE_CHECKING: from .llm_args import MoeConfig +LAUNCH_QUEUES_ENV = "CUDA_SCALE_LAUNCH_QUEUES" +_FROM_ENVIRON = object() + def configure_moe_launch_queues( moe_config: MoeConfig | None, env_overrides: dict[str, Any] | None = None, + *, + initial_setting: Any = _FROM_ENVIRON, ) -> dict[str, Any] | None: """Configure the per-layer dynamic EPLB launch queue before CUDA setup. @@ -27,6 +32,11 @@ def configure_moe_launch_queues( CUDA has no supported operation to resize existing launch queues here. The late-init guard detects PyTorch initialization; external CUDA users must still arrange for this setting before initializing their contexts. + + ``initial_setting`` is this process's ``CUDA_SCALE_LAUNCH_QUEUES`` before + the caller replayed ``env_overrides``: an existing CUDA context was created + with that value, not with the replayed one. It defaults to the current + environment. """ if moe_config is not None: moe_config.resolve_load_balancer_compatibility() @@ -34,9 +44,11 @@ def configure_moe_launch_queues( if getattr(load_balancer, "mode", None) != "per_layer_dynamic": return env_overrides - name = "CUDA_SCALE_LAUNCH_QUEUES" + name = LAUNCH_QUEUES_ENV + if initial_setting is _FROM_ENVIRON: + initial_setting = os.environ.get(name) torch = sys.modules.get("torch") - if os.environ.get(name) != "4x" and torch is not None and torch.cuda.is_initialized(): + if initial_setting != "4x" and torch is not None and torch.cuda.is_initialized(): raise RuntimeError( "Per-layer dynamic EPLB requires CUDA_SCALE_LAUNCH_QUEUES=4x " "before CUDA initialization. Restart the process with that " diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index c9bb16835e05..486cade00086 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -55,7 +55,7 @@ from ..logger import logger from ..sampling_params import LogitsProcessor, SamplingParams from ..scheduling_params import SchedulingParams -from ._load_balance_env import configure_moe_launch_queues +from ._load_balance_env import LAUNCH_QUEUES_ENV, configure_moe_launch_queues from .llm_args import (ENCODER_RUNNER_MANAGED_INPUTS, TORCH_LLMARGS_EXPLICIT_DOCSTRING, TORCH_LLMARGS_REMOVED_ARGS, ExecutorMemoryType, @@ -385,6 +385,7 @@ def __init__(self, try: env_overrides = kwargs.get("env_overrides", None) + launch_queues = os.environ.get(LAUNCH_QUEUES_ENV) # Rebalance validation and launch-queue setup read process state. # Apply caller overrides before either one observes that state. self._process_env_overrides(env_overrides) @@ -395,7 +396,7 @@ def __init__(self, moe_config = MoeConfig(**moe_config) kwargs["moe_config"] = moe_config queue_overrides = configure_moe_launch_queues( - moe_config, env_overrides) + moe_config, env_overrides, initial_setting=launch_queues) if queue_overrides is not env_overrides: kwargs["env_overrides"] = queue_overrides env_overrides = queue_overrides diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index d71c20718973..b752a228f739 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -1080,6 +1080,16 @@ def config(mode="per_layer_dynamic"): module.configure_moe_launch_queues(config(), original) assert original == {"KEEP": "value"} assert _QUEUE not in test_environ + + # A replayed override is too late for a context that already exists: + # the check reads the setting the process had before the replay. + test_environ[_QUEUE] = "4x" + with pytest.raises(RuntimeError, match="before CUDA initialization"): + module.configure_moe_launch_queues(config(), original, initial_setting=None) + assert module.configure_moe_launch_queues(config(), original, initial_setting="4x") == { + "KEEP": "value", + _QUEUE: "4x", + } assert os.environ.get(_QUEUE) == process_queue @@ -1118,9 +1128,10 @@ def test_launch_queue_setup_sees_env_overrides_before_runtime_initialization(mon user_overrides = {"USER_SETTING": "1"} configured = {**user_overrides, _QUEUE: "4x"} - def configure_moe_launch_queues(moe_config, env_overrides): - # The override must already be visible when the launch queue is configured. - events.append(("configure", environ.get("USER_SETTING"))) + def configure_moe_launch_queues(moe_config, env_overrides, *, initial_setting): + # The override must already be visible when the launch queue is configured, + # while the late-init check sees the setting from before the replay. + events.append(("configure", environ.get("USER_SETTING"), initial_setting)) return {**(env_overrides or {}), _QUEUE: "4x"} def runtime_init(name): @@ -1142,11 +1153,12 @@ def worker_args(): backend="pytorch", env_overrides=dict(user_overrides), moe_config=object() ) - expected = [("env", "USER_SETTING", "1"), ("configure", "1")] + expected = [("env", "USER_SETTING", "1"), ("configure", "1", None)] common = { "os": fake_os, "logger": logger, "configure_moe_launch_queues": configure_moe_launch_queues, + "LAUNCH_QUEUES_ENV": _QUEUE, } # LLM API: argument validation is the first step that may probe CUDA. @@ -2779,3 +2791,24 @@ def barrier(self): assert manager.shutdown(local_safe=True) is False assert calls == [("allgather", False)] + + +def test_standard_load_balancer_finalizer_releases_it_as_on_main(): + manager_cls = _extract_class( + _TORCH_ROOT / "moe/fused_moe/moe_load_balancer.py", + "MoeLoadBalancer", + {"__del__"}, + {}, + ) + calls = [] + manager = manager_cls() + manager.shutdown = lambda: calls.append("shutdown") + + manager.is_shutdown = True + manager.__del__() + assert calls == [] + + manager.is_shutdown = False + manager.__del__() + assert calls == ["shutdown"] + manager.is_shutdown = True From b4d9ffbf6b7fbfbcd7f0a8d50694bb4288e27138 Mon Sep 17 00:00:00 2001 From: jintaop Date: Mon, 26 Oct 2026 08:46:17 -0700 Subject: [PATCH 43/59] [None][fix] Keep MegaMoE per-expert scalar defaults when binding the EPLB arena Binding MegaMoE's live planes onto the per-layer dynamic EPLB arena rebinds each Parameter to zero-initialized arena pages without copying its create-time value. Only the helper rows of fc1_norm_const were reset to 1.0, so fc31_alpha, fc2_alpha and fc1_norm_const stayed 0 for every resident expert whose checkpoint does not set them, and that expert's routed output became exactly 0 instead of matching the run without the load balancer. This masked a wrong layer. DeepSeek-V4-Pro-NVFP4 ships MXFP4 MTP experts without weight_scale_2 or input_scale; they are loaded as NVFP4 on main. On GB200 (DEP8, MTPDecodingConfig(max_draft_len=1), CUDA graphs off) the MTP acceptance length was 1.84-1.92 with per-layer dynamic EPLB, where layer 61's routed output was zeroed, and 1.01 without it, in both run orders. Restore 1.0 on every row of the three planes so both modes compute the same layer. Signed-off-by: jintaop --- .../_torch/moe/fused_moe/quantization.py | 11 +-- .../_torch/moe/test_dynamic_eplb_contracts.py | 72 +++++++++++++++++++ 2 files changed, 79 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/quantization.py b/tensorrt_llm/_torch/moe/fused_moe/quantization.py index 1607bd2af75f..fa9ca966451a 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/moe/fused_moe/quantization.py @@ -4428,10 +4428,13 @@ def _bind_live_weight_planes(self, module: torch.nn.Module, arena) -> None: ) from error if rebound: self.setup_quant_scales(module) - norm_const = getattr(module, "fc1_norm_const", None) - if (norm_const is not None - and arena.home_experts < norm_const.data.shape[0]): - norm_const.data[int(arena.home_experts):].fill_(1.0) + # Arena pages start zeroed. Restore the create-time value of the + # per-expert scalars on every row, so an expert whose checkpoint leaves + # them unset computes as it does without the load balancer. + for name in ("fc31_alpha", "fc2_alpha", "fc1_norm_const"): + plane = getattr(module, name, None) + if plane is not None: + plane.data.fill_(1.0) self.assert_live_weight_aliases(module, arena) # ----------------------------------------------------------------- diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index b752a228f739..be190575d3a9 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -2812,3 +2812,75 @@ def test_standard_load_balancer_finalizer_releases_it_as_on_main(): manager.__del__() assert calls == ["shutdown"] manager.is_shutdown = True + + +class _FakePlane: + """Stand-in for an arena alias: records fills, passes the layout checks.""" + + is_cuda = True + dtype = "float32" + + def __init__(self, name, address, fills): + self.name = name + self.shape = (6,) + self._address = address + self._fills = fills + + def stride(self): + return (1,) + + def data_ptr(self): + return self._address + + def fill_(self, value): + self._fills.append((self.name, value)) + + +def test_live_plane_binding_restores_per_expert_scalar_defaults(): + names = ( + "mega_fc1_weight", + "mega_fc1_weight_sf", + "mega_fc2_weight", + "mega_fc2_weight_sf", + "fc31_alpha", + "fc2_alpha", + "fc1_norm_const", + ) + fills = [] + aliases = [_FakePlane(name, 0x1000 + index, fills) for index, name in enumerate(names)] + arena = SimpleNamespace( + bundle="bundle", + plane_names=names, + local_plane_views=aliases, + tekit_alias_views=aliases, + home_experts=4, + ) + method_cls = _extract_class( + _TORCH_ROOT / "moe/fused_moe/quantization.py", + "NVFP4MegaMoECuteDslMethod", + {"_bind_live_weight_planes"}, + {"torch": SimpleNamespace(nn=SimpleNamespace(Parameter=None))}, + ) + method = method_cls() + method.live_weight_plane_spec = lambda module: SimpleNamespace( + planes=[SimpleNamespace(name=name) for name in names] + ) + arena.bundle = method.live_weight_plane_spec(None) + method.assert_live_weight_aliases = lambda module, bound: None + module = SimpleNamespace( + **{ + name: SimpleNamespace( + is_meta=False, shape=(6,), stride=lambda: (1,), dtype="float32", data=None + ) + for name in names + } + ) + + method._bind_live_weight_planes(module, arena) + + # The arena starts zeroed; an expert the checkpoint leaves unset must keep + # the create-time scalars (1.0) instead of silently computing zeros. + assert sorted(fills) == sorted( + [("fc31_alpha", 1.0), ("fc2_alpha", 1.0), ("fc1_norm_const", 1.0)] + ) + assert all(getattr(module, name).data is alias for name, alias in zip(names, aliases)) From fae35213b35e229e33bc84a6bb4e4f6edfd2f7bf Mon Sep 17 00:00:00 2001 From: jintaop Date: Mon, 26 Oct 2026 08:46:34 -0700 Subject: [PATCH 44/59] [None][doc] State which draft MoE layers per-layer dynamic EPLB balances Every MoE layer that runs on MEGAMOE_CUTEDSL, including one-model draft layers, registers with the per-layer dynamic load balancer. Document that DSpark drafters in an NVFP4 target checkpoint are balanced this way, and that MTP on DeepSeek-V4-Pro-NVFP4 is not supported because its MTP experts are MXFP4. Validation of PR head b5f61c3c6: - DSpark, DeepSeek-V4-Pro-nvfp4-DSpark on VR200, DEP4, MEGAMOE_CUTEDSL, max_draft_len 5, CUDA graphs off, 300 GSM8K samples: per-layer dynamic EPLB (num_slots 396) 96.83 with acceptance length 4.45, EPLB off 96.67 with 4.44. The three draft stages registered as per-layer dynamic layers 61-63 next to the 61 target layers, and both runs exited cleanly. - MTP, DeepSeek-V4-Pro-NVFP4 on GB200, DEP8: layer 61 builds as NVFP4 MegaMoE and registers with the load balancer, but its MXFP4 checkpoint experts are loaded as NVFP4; drafts are almost all rejected without the load balancer (acceptance length 1.01). Signed-off-by: jintaop --- examples/wide_ep/ep_load_balancer/README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/wide_ep/ep_load_balancer/README.md b/examples/wide_ep/ep_load_balancer/README.md index 36aa87b31230..0f1883ea4b62 100644 --- a/examples/wide_ep/ep_load_balancer/README.md +++ b/examples/wide_ep/ep_load_balancer/README.md @@ -311,5 +311,6 @@ Per-layer dynamic EPLB has the following constraints: * **Launch queues**: CUDA must be initialized with `CUDA_SCALE_LAUNCH_QUEUES=4x`. The LLM API sets it automatically, in the main process and in every worker, before CUDA is initialized. If the process has already initialized CUDA when the LLM is created, construction fails; export `CUDA_SCALE_LAUNCH_QUEUES=4x` before starting such a process. * **`auxiliary_sms`**: the number of SMs of the kernel that runs the HALO-Q scheduler and then the copy, one CTA per SM. It defaults to 4, must be between 1 and 8, and must be lower than the GPU's SM count. Input quantization and, on the SM100 family, DeepSeek-V4's fused shared-expert kernel leave these SMs free, so planning and copying run alongside them. * **`helper_slot_mode`**: `per_layer_resident` (the default) gives every MoE layer its own helper slots and skips copying experts that they still hold from the previous forward pass. It costs one helper set per MoE layer, about 6.4 GiB per GPU for DeepSeek-V4-Pro with 3 helper slots per rank. `shared_banks` has all MoE layers alternate between two shared helper banks and copies every helper expert in every forward pass. +* **Speculative decoding**: a one-model draft's MoE layers follow the same rule as the target's: every MoE layer that runs on `MEGAMOE_CUTEDSL` registers as one more per-layer dynamic layer after the target layers, is planned and load-balanced like them, and, with `per_layer_resident`, gets its own helper slots. DSpark drafters shipped inside an NVFP4 target checkpoint, such as `DeepSeek-V4-Pro-nvfp4-DSpark`, do so with NVFP4 draft experts. A draft layer only computes correctly if its checkpoint experts are NVFP4: the `mtp.*` routed experts of the `DeepSeek-V4-Pro-NVFP4` checkpoint are MXFP4, which `MEGAMOE_CUTEDSL` does not serve, so MTP on that checkpoint is not supported with or without this load balancer. * **DWDP**: not supported; `dwdp_config` must keep `dwdp_size` at 1. * **Fields**: only `mode`, `num_slots`, `auxiliary_sms` and `helper_slot_mode`. `num_slots` is required, must be greater than the expert count, and must be divisible by the EP size. `layer_updates_per_iter` and `initial_global_assignments` are `standard`-only and rejected in this mode. From 161db80e62de6873579a38c80b611e5f7d8802d6 Mon Sep 17 00:00:00 2001 From: jintaop Date: Mon, 26 Oct 2026 21:25:59 -0700 Subject: [PATCH 45/59] [None][perf] Skip no-op hooks and repeated checks on the per-layer dynamic EPLB path Per-layer dynamic EPLB runs its host bookkeeping once per MoE layer and forward, and on Grace in-situ Python is several times slower than in a hot loop, so every call on this path counts. Outputs are unchanged: - the fused-comm MoE scheduler no longer calls the CPU-stage and statistics hooks for a layer that produces a replica plan; they are no-ops for it. It reads the plan from the layer directly, and the now unused _load_balancer_replica_plan helper is removed; - plan_finish() no longer re-checks the MAIN-stream owner: it queues no device work, and plan_schedule() and the route wait check it; - the route check compares the routes' device index instead of building torch.device objects, and the route wait resolves its driver constants once; - the layer handle resolves its READY-flag view once, at bind; - the reserved-SM FP4 quantize passes reserved_sms positionally. Signed-off-by: jintaop --- .../_torch/moe/fused_moe/interface.py | 6 --- .../fused_moe/mega_moe/mega_moe_cute_dsl.py | 8 +-- .../_torch/moe/fused_moe/moe_scheduler.py | 51 +++++++++++-------- .../fused_moe/per_layer_dynamic_eplb/layer.py | 6 ++- .../per_layer_dynamic_eplb/planner.py | 20 +++++--- .../_torch/moe/test_dynamic_eplb_contracts.py | 36 ++++++++++--- tests/unittest/_torch/moe/test_moe_backend.py | 2 +- 7 files changed, 78 insertions(+), 51 deletions(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/interface.py b/tensorrt_llm/_torch/moe/fused_moe/interface.py index b13c06dbd262..6d7754a74586 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/interface.py +++ b/tensorrt_llm/_torch/moe/fused_moe/interface.py @@ -773,12 +773,6 @@ def _load_balancer_uses_replica_plan(self) -> bool: return bool(self.layer_load_balancer and self.layer_load_balancer.produces_replica_plan) - def _load_balancer_replica_plan(self): - """Return the plan produced by the most recent route, if this mode has one.""" - if not self._load_balancer_uses_replica_plan(): - return None - return self.layer_load_balancer.require_replica_plan() - def _load_balancer_start_set_cpu_stage(self, is_last_call: bool): """Start CPU stage in load balancer.""" # A per-layer dynamic EPLB layer plans every chunk and releases that diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py index 1329c405c95b..3ca2c6889890 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py @@ -1255,13 +1255,9 @@ def quantize_input( # Leave a replica plan's SMs to its concurrent HALO-Q/TMA kernel. The # default overload remains unchanged without a plan. if reserved_sms > 0 and _fp4_quantize_has_sm_budget(): + # Positional: a keyword argument costs the op dispatch about a microsecond. x_fp4, x_sf = torch.ops.trtllm.fp4_quantize.sm_budget( - x_bf16, - self.fc31_input_scale, - 16, - False, - False, - reserved_sms=reserved_sms, + x_bf16, self.fc31_input_scale, 16, False, False, reserved_sms ) else: x_fp4, x_sf = torch.ops.trtllm.fp4_quantize( diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py index ac99e08f2308..f07aa1e87856 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py @@ -1492,7 +1492,10 @@ def _forward_chunk( quantize -> ``run_moe`` -> ``start_set_cpu_stage`` -> ``done_set_cpu_stage``. ``start/done_set_cpu_stage`` are placed AFTER ``run_moe``; otherwise dynamic-EPLB weight migration would - race with the fused kernel using those weights. + race with the fused kernel using those weights. A per-layer dynamic + EPLB layer has no CPU stage and no host statistics: it only gets + ``route`` and the post-``run_moe`` ``start_set_cpu_stage``, which + releases its replica plan. """ moe = self.moe assert not moe.apply_router_weight_on_input, ( @@ -1521,10 +1524,15 @@ def _forward_chunk( _leading_rows(input_ids, num_tokens) if input_ids is not None else None ) + # Per-layer dynamic EPLB plans on the GPU; the CPU-stage and + # statistics hooks below are no-ops for it, so it skips them. + uses_replica_plan = moe._load_balancer_uses_replica_plan() + # ----- EPLB: drain previous CPU rebalance ----- # Static EPLB early-returns inside the helper; only the dynamic # balancer actually waits. - moe._load_balancer_start_wait_gpu_stage(is_first_call) + if not uses_replica_plan: + moe._load_balancer_start_wait_gpu_stage(is_first_call) # ----- routing ----- # int32 matches the EPLB stats kernel contract used by the external-comm @@ -1567,27 +1575,27 @@ def _forward_chunk( # Per-layer dynamic EPLB plans from the routes that the layer-wise # calibrator records or replays, so they are recorded before the route; # otherwise the recorded routes are the slots routed to. - uses_replica_plan = moe._load_balancer_uses_replica_plan() - if moe.layer_load_balancer: - moe._load_balancer_done_wait_gpu_stage(is_first_call) - # ignore_allreduce=False: the fused kernel has no side channel - # for an external stats gather. The base helper runs its own - # EP-wide AllReduce, gated to is_last_call=True. - moe._load_balancer_update_statistic( - token_selected_experts, - is_first_call, - is_last_call, - ignore_allreduce=False, - ) - if uses_replica_plan: - token_selected_experts = _record_routes(moe, token_selected_experts) + if uses_replica_plan: + token_selected_experts = _record_routes(moe, token_selected_experts) token_selected_slots = moe._load_balancer_route(token_selected_experts, moe.use_dp) + replica_plan = moe.layer_load_balancer.require_replica_plan() else: - token_selected_slots = token_selected_experts - - if not uses_replica_plan: + if moe.layer_load_balancer: + moe._load_balancer_done_wait_gpu_stage(is_first_call) + # ignore_allreduce=False: the fused kernel has no side channel + # for an external stats gather. The base helper runs its own + # EP-wide AllReduce, gated to is_last_call=True. + moe._load_balancer_update_statistic( + token_selected_experts, + is_first_call, + is_last_call, + ignore_allreduce=False, + ) + token_selected_slots = moe._load_balancer_route(token_selected_experts, moe.use_dp) + else: + token_selected_slots = token_selected_experts token_selected_slots = _record_routes(moe, token_selected_slots) - replica_plan = moe._load_balancer_replica_plan() if uses_replica_plan else None + replica_plan = None # ----- quantize / prepare ----- if getattr(moe.backend, "supports_fused_prepare", lambda: False)(): @@ -1639,7 +1647,8 @@ def _forward_chunk( # start_set before run_moe would let dynamic-EPLB migration race the # kernel. moe._load_balancer_start_set_cpu_stage(is_last_call) - moe._load_balancer_done_set_cpu_stage(is_last_call) + if not uses_replica_plan: + moe._load_balancer_done_set_cpu_stage(is_last_call) return out diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/layer.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/layer.py index a1d645109f06..69fd6b628833 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/layer.py +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/layer.py @@ -65,6 +65,8 @@ def __init__( self.repeated_count = int(repeated_count) self._arena: Optional[Any] = None + # The arena's READY terminal view, resolved once the runtime is bound. + self._ready_flags: Optional[torch.Tensor] = None self._scheduler_group: Optional[Any] = None self._bound_backend_ref: Optional[weakref.ReferenceType[Any]] = None self._pending_replica_plan: Optional[Any] = None @@ -170,6 +172,7 @@ def bind_runtime(self, *, arena: Any, scheduler_group: Any) -> None: raise ValueError("scheduler H/S geometry does not match the EPLB layer") self._arena = arena + self._ready_flags = arena.terminal_flags_tensor() self._scheduler_group = scheduler_group self.manager._register_runtime(self) @@ -317,7 +320,7 @@ def route( plan = MoEReplicaPlan( resident_experts_per_rank=self.resident_experts_per_rank, compute_slots_per_rank=self.compute_slots_per_rank, - ready_flags=self._arena.terminal_flags_tensor(), + ready_flags=self._ready_flags, ready_generation=generation, reserved_sms=self.manager.auxiliary_sms, _wait_for_routes=self._scheduler_group.wait_for_routes, @@ -356,4 +359,5 @@ def _clear_runtime(self) -> None: self._pending_replica_plan = None self._scheduler_group = None self._arena = None + self._ready_flags = None self._bound_backend_ref = None diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py index 72f0c8f186c1..f10de29e9f51 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py @@ -205,6 +205,9 @@ def __init__( self.scheduler_stream = self.copy_stream self._stream_handle = int(self.copy_stream.cuda_stream) self._driver = driver + # The route wait runs once per layer; resolve its driver constants once. + self._wait_value_geq = driver.CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_GEQ + self._cuda_success = driver.CUresult.CUDA_SUCCESS error, cu_device = driver.cuDeviceGet(self.device) supported = 0 if error == driver.CUresult.CUDA_SUCCESS: @@ -309,12 +312,9 @@ def _check_owner(self) -> None: def _wait_route_epoch(self, stream_handle: int) -> None: (error,) = self._driver.cuStreamWaitValue64( - stream_handle, - self._route_epoch_ptr, - self._route_target, - self._driver.CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_GEQ, + stream_handle, self._route_epoch_ptr, self._route_target, self._wait_value_geq ) - if error != self._driver.CUresult.CUDA_SUCCESS: + if error != self._cuda_success: raise RuntimeError(f"cuStreamWaitValue64(rebalance) failed: {error!r}") def plan_schedule(self, logical_expert_ids: torch.Tensor) -> Tuple[Any, int]: @@ -341,7 +341,8 @@ def plan_schedule(self, logical_expert_ids: torch.Tensor) -> Tuple[Any, int]: or logical_expert_ids.shape[1] != self.topk or logical_expert_ids.dtype != torch.int32 or not logical_expert_ids.is_contiguous() - or logical_expert_ids.device != self.scheduler.device + # -1 for a CPU tensor; cheaper than building and comparing torch.device objects. + or logical_expert_ids.get_device() != self.device ): raise ValueError("Routes must be contiguous CUDA int32 [tokens, topk] on this device") num_tokens = int(logical_expert_ids.shape[0]) @@ -374,8 +375,11 @@ def plan_schedule(self, logical_expert_ids: torch.Tensor) -> Tuple[Any, int]: def plan_finish( self, part: Tuple[Any, int], *, defer_wait: bool = False ) -> Tuple[torch.Tensor, int]: - """Lend routes until finish(); optionally defer their device dependency.""" - self._check_owner() + """Lend routes until finish(); optionally defer their device dependency. + + It queues no device work itself: plan_schedule() checked the owner of + this pending plan, and the route wait checks it again. + """ if part is not self._plan_part or part is None: raise RuntimeError("plan_finish requires this group's pending plan") outputs, num_tokens = part diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index be190575d3a9..9c7c490a8509 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -1005,6 +1005,9 @@ def __init__( def is_contiguous(self): return self._contiguous + def get_device(self): + return int(self.device.split(":")[1]) if self.device.startswith("cuda:") else -1 + def record_stream(self, stream): self.trace.append(("lifetime", stream.cuda_stream)) @@ -1302,6 +1305,8 @@ def driver_call(*entry): "wait_value", stream, address, value, flags ), ) + group._wait_value_geq = group._driver.CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_GEQ + group._cuda_success = group._driver.CUresult.CUDA_SUCCESS outputs = SimpleNamespace(physical_slot_ids=_Routes(8192, trace)) launch_bounds = [] @@ -1851,8 +1856,8 @@ def fp4_quantize(*args): calls.append(("default", *args[1:])) return "x_fp4", Activation(3, 4) - def sm_budget(*args, reserved_sms): - calls.append(("sm_budget", *args[1:], reserved_sms)) + def sm_budget(*args): + calls.append(("sm_budget", *args[1:])) return "x_fp4", Activation(3, 4) if has_sm_budget: @@ -2411,7 +2416,14 @@ def maybe_collect_or_replay_slots(num_slots, routes): scheduler_cls = _extract_class(source, "FusedCommMoEScheduler", {"_forward_chunk"}, namespace) plan = SimpleNamespace(reserved_sms=4) - layer = None if mode == "off" else SimpleNamespace(produces_replica_plan=mode != "standard") + layer = ( + None + if mode == "off" + else SimpleNamespace( + produces_replica_plan=mode != "standard", + require_replica_plan=lambda: plan, + ) + ) def route(routes, use_dp): trace.append(("route", routes.name)) @@ -2435,14 +2447,13 @@ def run_moe(ctx): ) ), _maybe_get_perfect_router_expert_ids=lambda routes: routes, - _load_balancer_start_wait_gpu_stage=lambda is_first_call: None, - _load_balancer_done_wait_gpu_stage=lambda is_first_call: None, + _load_balancer_start_wait_gpu_stage=lambda is_first_call: trace.append("start_wait_gpu"), + _load_balancer_done_wait_gpu_stage=lambda is_first_call: trace.append("done_wait_gpu"), _load_balancer_update_statistic=lambda routes, *args, **kwargs: trace.append( ("update_statistic", routes.name) ), _load_balancer_uses_replica_plan=lambda: mode == "per_layer_dynamic", _load_balancer_route=route, - _load_balancer_replica_plan=lambda: plan if mode == "per_layer_dynamic" else None, _load_balancer_start_set_cpu_stage=lambda is_last_call: trace.append("start_set_cpu"), _load_balancer_done_set_cpu_stage=lambda is_last_call: trace.append("done_set_cpu"), backend=SimpleNamespace( @@ -2468,32 +2479,41 @@ def run_moe(ctx): if mode == "per_layer_dynamic": # The planner consumes the logical routes the calibrator records or # replays, and quantization leaves the plan's SMs to its copy kernel. + # The layer has no CPU stage or host statistics, so only the post-MoE + # hook that releases its plan runs. expected = [ - ("update_statistic", "logical"), ("statistic", "logical"), ("calibrator", "logical"), ("route", "replayed logical"), ("quantize", {"reserved_sms": 4}), ("bucket", 6), ("run_moe", "physical", plan), + "start_set_cpu", ] elif mode == "standard": expected = [ + "start_wait_gpu", + "done_wait_gpu", ("update_statistic", "logical"), ("route", "logical"), ("statistic", "slots"), ("calibrator", "slots"), *tail, ("run_moe", "replayed slots", None), + "start_set_cpu", + "done_set_cpu", ] else: expected = [ + "start_wait_gpu", ("statistic", "logical"), ("calibrator", "logical"), *tail, ("run_moe", "replayed logical", None), + "start_set_cpu", + "done_set_cpu", ] - assert trace == expected + ["start_set_cpu", "done_set_cpu"] + assert trace == expected def test_base_worker_finalizes_rank_follower_pyexecutor(): diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index ca01bc542577..59818abdf51d 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -2473,7 +2473,7 @@ def fp4_quantize(x, scale, sf_vec_size, sf_use_ue8m0, is_sf_swizzled): calls.append(("default", sf_vec_size, sf_use_ue8m0, is_sf_swizzled)) return outputs(x) - def sm_budget(x, scale, sf_vec_size, sf_use_ue8m0, is_sf_swizzled, *, reserved_sms): + def sm_budget(x, scale, sf_vec_size, sf_use_ue8m0, is_sf_swizzled, reserved_sms): calls.append(("sm_budget", sf_vec_size, sf_use_ue8m0, is_sf_swizzled, reserved_sms)) return outputs(x) From 91d8c0b26becab2c65d59dac1c655400703da685 Mon Sep 17 00:00:00 2001 From: jintaop Date: Mon, 26 Oct 2026 21:26:34 -0700 Subject: [PATCH 46/59] [None][perf] Let the next fused launch order COPY after the per-layer EPLB consumer After each MegaMoE launch, finish() released the plan with a native call that recorded a consumer event on MAIN and made COPY wait on it. That wait is implied: every fused HALO-Q + TMA launch already records an event on the current stream and makes COPY wait on it before it starts, and the next launch of any layer is submitted on the same bound MAIN stream after that MegaMoE. finish() still checks that MAIN is the current stream, and a caller-recorded completion event is still waited for on COPY. This removes one torch op and two CUDA runtime calls per MoE layer. The ordering is unchanged: COPY reaches each layer's next HALO-Q exchange and helper copy only after the MAIN work queued before it, which includes the previous consumer of that layer's helper slots. Signed-off-by: jintaop --- .../fused_moe/per_layer_dynamic_eplb/planner.py | 17 ++++++++++------- .../megamoe_scheduler/sami/hierarchical.py | 15 ++++++++------- .../_torch/moe/test_dynamic_eplb_contracts.py | 17 +++++------------ 3 files changed, 23 insertions(+), 26 deletions(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py index f10de29e9f51..237c89e85344 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py @@ -11,8 +11,10 @@ protects helper reads while the copy finishes concurrently. For each rank and generation g: - MAIN: MegaMoE(g) -> consumer_done(g) - COPY: wait consumer_done(g) -> [HALO-Q(g+1) EP exchange -> TMA(g+1)] + MAIN: MegaMoE(g) -> ... -> routes(g+1) -> input_ready(g+1) + COPY: wait input_ready(g+1) -> [HALO-Q(g+1) EP exchange -> TMA(g+1)] +Every launch orders COPY after the MAIN work queued before it, which includes +the previous consumer MegaMoE(g), so finish() queues no device work of its own. All ranks enter the bounded EP exchange after their own previous consumer, so no next-generation overwrite can precede any previous-generation consumer. No host collective or auxiliary copy submitter is needed in the serving path. @@ -76,9 +78,9 @@ def cycles(self) -> int: class _LiveBankLeaseProvider: """Authorize ordered reuse, not a claim of completed GPU work on the host. - release_generation_after() enqueues the local consumer wait on COPY. - prove() is legal right before the next fused HALO-Q + TMA launch on that - same stream. The launch's bounded all-rank count exchange precedes its plan, + release_generation_after() registers the local consumer; the next fused + HALO-Q + TMA launch orders COPY after it on MAIN. prove() is legal right + before that launch. The launch's bounded all-rank count exchange precedes its plan, and its copy follows the plan in the same CTAs. Therefore a delayed peer consumer blocks every rank's overwrite without a host-side barrier. """ @@ -264,7 +266,7 @@ def __init__( ) self.broadcaster.bind_scheduler(self.scheduler, self.scheduler_stream, self._route_epoch) - # A prior consumer event precedes the next HALO-Q EP rendezvous. + # The previous consumer precedes the next HALO-Q EP rendezvous. self.lease = _LiveBankLeaseProvider(self) self.broadcaster.bind_generation_reuse_authority(self.lease) @@ -419,10 +421,11 @@ def discard_plan(self, part: Tuple[Any, int]) -> None: self.finish() def finish(self, completion_event: torch.cuda.Event | None = None) -> None: - """Queue the previous consumer wait before any next-generation producer. + """Release this generation to the next-generation producer. The consumer is ``completion_event``, or without it the MAIN work queued so far: the MegaMoE launch that read this generation, if any. + The next fused launch, submitted on MAIN, orders COPY after it. """ self._check_owner() if ( diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py index 6f4a8975fa47..9f3e4befdb37 100644 --- a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py @@ -587,10 +587,13 @@ def release_generation_after( """Order the copy stream after this rank's live-bank consumer. The consumer is ``completion_event``, or without it the work already - queued on the current stream. This establishes only local stream - ordering. The bound reuse authority must separately attest its - bounded all-rank barrier before another generation may overwrite the - single live bank. + queued on the current stream. That stream must be the one the next + fused launch follows: the launch orders the copy stream after the work + queued on its current stream (see ``submit_fused``), which then + includes this consumer, so no separate device wait is queued for it. + This establishes only local stream ordering. The bound reuse authority + must separately attest its bounded all-rank barrier before another + generation may overwrite the single live bank. """ if type(generation) is not int or generation <= 0: @@ -607,9 +610,7 @@ def release_generation_after( raise ValueError("generation is not the active live bank") if self._local_release_generation >= generation: raise RuntimeError("generation already has a local release event") - if completion_event is None: - torch.ops.trtllm.moe_rebalance_tma_release(self._handle) - else: + if completion_event is not None: check_cuda( self.cuda.cuStreamWaitEvent(self._stream, event_handle, 0), "cuStreamWaitEvent(live-bank consumer)", diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index 9c7c490a8509..38028fb91e3a 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -239,16 +239,8 @@ def test_hierarchical_fused_submit_launches_bound_scheduler_and_copy(): def submit_fused(handle, routes, route_target, spin_cycles): native.append(("launch", handle, routes, route_target, spin_cycles)) - def release(handle): - native.append(("release", handle)) - torch = SimpleNamespace( - ops=SimpleNamespace( - trtllm=SimpleNamespace( - moe_rebalance_tma_submit_fused=submit_fused, - moe_rebalance_tma_release=release, - ) - ), + ops=SimpleNamespace(trtllm=SimpleNamespace(moe_rebalance_tma_submit_fused=submit_fused)), ) ticket_cls = _extract_class( source, @@ -322,10 +314,11 @@ def revalidate(_authority): with pytest.raises(RuntimeError, match="local consumer release"): broadcast.mark_collective_reuse_safe(authority, 1) native.clear() - # Without an event, the consumer is the work queued on the current stream, - # which one native call orders the copy stream after. + # Without an event, the consumer is the work queued on the current stream. + # The next launch orders the copy stream after that stream, so the release + # queues no device work. broadcast.release_generation_after(1) - assert native == [("release", 4242)] + assert native == [] with pytest.raises(RuntimeError, match="already has a local release"): broadcast.release_generation_after(1) broadcast.mark_collective_reuse_safe(authority, 1) From f79d9c5623d25b05ee3c7a0c151caafe813324bb Mon Sep 17 00:00:00 2001 From: jintaop Date: Mon, 26 Oct 2026 21:46:44 -0700 Subject: [PATCH 47/59] [None][perf] Return the launched generation from the per-layer EPLB submit The SAMI broadcaster's fused submit built a frozen HierarchicalCopyTicket per MoE layer, and the planner read only its generation. It now returns the generation itself, and the unused ticket class is removed. The plan release calls the group's route wait directly instead of through the replica plan's callback. Outputs and ordering are unchanged. Signed-off-by: jintaop --- .../fused_moe/per_layer_dynamic_eplb/layer.py | 5 ++-- .../per_layer_dynamic_eplb/planner.py | 6 ++--- .../moe/kernels/megamoe_scheduler/__init__.py | 2 -- .../megamoe_scheduler/sami/__init__.py | 2 -- .../megamoe_scheduler/sami/hierarchical.py | 23 +++++-------------- .../_torch/moe/test_dynamic_eplb_contracts.py | 18 ++++----------- 6 files changed, 17 insertions(+), 39 deletions(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/layer.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/layer.py index 69fd6b628833..a81315445f82 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/layer.py +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/layer.py @@ -350,8 +350,9 @@ def finish_replica_plan( raise RuntimeError("attempted to finish a different replica plan") assert self._scheduler_group is not None # Also makes cleanup safe when input preparation failed before the - # backend reached its first route consumer. - pending.wait_for_routes() + # backend reached its first route consumer. This is the plan's route + # wait, called on the group directly. + self._scheduler_group.wait_for_routes() self._scheduler_group.finish(completion_event=completion_event) self._pending_replica_plan = None diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py index 237c89e85344..af66b06e05fa 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py @@ -360,16 +360,16 @@ def plan_schedule(self, logical_expert_ids: torch.Tensor) -> Tuple[Any, int]: self._routes_in_flight = logical_expert_ids # The launch follows the MAIN work queued so far, which wrote the # routes; no host synchronization is needed. - ticket = self.broadcaster.submit_fused( + generation = self.broadcaster.submit_fused( outputs, logical_expert_ids, route_target, spin_cycles=self._spin_bound.cycles(), ) - if int(ticket.generation) != self._scheduled_generation: + if generation != self._scheduled_generation: raise RuntimeError("TMA and HALO-Q generation mismatch") self._route_target = route_target - self._generation = int(ticket.generation) + self._generation = generation part = (outputs, num_tokens) self._plan_part = part return part diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/__init__.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/__init__.py index 1516768d758e..efa1c69eff92 100644 --- a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/__init__.py @@ -11,7 +11,6 @@ ) from .sami import ( BoundHierarchicalLiveWeightArena, - HierarchicalCopyTicket, HierarchicalLiveWeightArena, HierarchicalLiveWeightArenaProvider, HierarchicalSamiWeightBroadcast, @@ -35,7 +34,6 @@ "BundleLayout", "Comm", "LivePlaneView", - "HierarchicalCopyTicket", "HierarchicalLiveWeightArena", "HierarchicalLiveWeightArenaProvider", "HierarchicalSamiWeightBroadcast", diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/__init__.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/__init__.py index 851fe65a18f0..ae2b7d7c8f7b 100644 --- a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/__init__.py @@ -14,7 +14,6 @@ from .geometry import BundleLayout, Plane, PlaneViewLayout from .hierarchical import ( BoundHierarchicalLiveWeightArena, - HierarchicalCopyTicket, HierarchicalLiveWeightArena, HierarchicalLiveWeightArenaProvider, HierarchicalSamiWeightBroadcast, @@ -26,7 +25,6 @@ "BoundLiveWeightArena", "BundleLayout", "Comm", - "HierarchicalCopyTicket", "HierarchicalLiveWeightArena", "HierarchicalLiveWeightArenaProvider", "HierarchicalSamiWeightBroadcast", diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py index 9f3e4befdb37..c4159e593d7b 100644 --- a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py @@ -182,13 +182,6 @@ def _event_handle(event: object | None) -> int: return int(getattr(event, "cuda_event", event)) -@dataclass(frozen=True) -class HierarchicalCopyTicket: - generation: int - terminal_flags: torch.Tensor = field(repr=False) - command_count: int - - class HierarchicalSamiWeightBroadcast: """One-rank HALO-Q copy endpoint with a single C hot-path context.""" @@ -658,7 +651,7 @@ def submit_fused( route_target: int, *, spin_cycles: int, - ) -> HierarchicalCopyTicket: + ) -> int: """Enqueue the bound scheduler and this GPU-direct copy as one kernel. The kernel runs on the bound stream after the work already queued on @@ -666,7 +659,8 @@ def submit_fused( adds one to the bound route epoch once its routes are written, which makes it reach ``route_target``; the copy and its READY publication follow in the same CTAs. ``spin_cycles`` bounds the launch's - cross-rank waits and its wait for the route epoch. + cross-rank waits and its wait for the route epoch. Returns the + launched generation. """ with self._lock: @@ -685,13 +679,9 @@ def submit_fused( except BaseException: self._handoff_aborted = True raise - ticket = HierarchicalCopyTicket( - generation=self._generation, - terminal_flags=self.terminal_flags, - command_count=1, - ) - self._generation += 1 - return ticket + generation = self._generation + self._generation = generation + 1 + return generation def close(self) -> None: """Release the native copy context after all device work is drained.""" @@ -719,7 +709,6 @@ def last_mode_counts(self) -> tuple[int, int, int]: __all__ = [ "BoundHierarchicalLiveWeightArena", - "HierarchicalCopyTicket", "HierarchicalLiveWeightArena", "HierarchicalLiveWeightArenaProvider", "HierarchicalSamiWeightBroadcast", diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index 38028fb91e3a..49afd1b25e26 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -242,12 +242,6 @@ def submit_fused(handle, routes, route_target, spin_cycles): torch = SimpleNamespace( ops=SimpleNamespace(trtllm=SimpleNamespace(moe_rebalance_tma_submit_fused=submit_fused)), ) - ticket_cls = _extract_class( - source, - "HierarchicalCopyTicket", - None, - {"dataclass": dataclasses.dataclass, "field": dataclasses.field}, - ) event_handle = _extract_function(source, "_event_handle", {}) broadcast_cls = _extract_class( source, @@ -260,7 +254,6 @@ def submit_fused(handle, routes, route_target, spin_cycles): }, { "torch": torch, - "HierarchicalCopyTicket": ticket_cls, "_event_handle": event_handle, "check_cuda": lambda result, operation: result, }, @@ -295,12 +288,12 @@ def revalidate(_authority): cuStreamWaitEvent=lambda stream, event, flags: native.append(("wait", stream, event)) ) - ticket = broadcast.submit_fused(outputs, "routes", 8, spin_cycles=12345) + generation = broadcast.submit_fused(outputs, "routes", 8, spin_cycles=12345) # One native call orders the launch after the current stream and enqueues # it on the bound copy stream; Python switches no stream. assert native == [("launch", 4242, "routes", 8, 12345)] - assert (ticket.generation, ticket.command_count) == (1, 1) + assert generation == 1 assert (broadcast._active_generation, broadcast._generation) == (1, 2) # The bank stays leased until the consumer of generation 1 is proven done. with pytest.raises(RuntimeError, match="still leased"): @@ -323,8 +316,7 @@ def revalidate(_authority): broadcast.release_generation_after(1) broadcast.mark_collective_reuse_safe(authority, 1) assert broadcast._collective_safe_generation == 1 - ticket = broadcast.submit_fused(outputs, "routes", 16, spin_cycles=12345) - assert ticket.generation == 2 + assert broadcast.submit_fused(outputs, "routes", 16, spin_cycles=12345) == 2 # A caller-recorded consumer event is waited for on the copy stream instead. native.clear() @@ -1309,7 +1301,7 @@ def submit_fused(result, routes, target, *, spin_cycles): generation = group._scheduled_generation trace.append(("FUSED", routes.shape[0], target, generation)) launch_bounds.append(spin_cycles) - return SimpleNamespace(generation=generation) + return generation group.scheduler = SimpleNamespace(device="cuda:0", outputs=outputs) group.broadcaster = SimpleNamespace( @@ -1369,7 +1361,7 @@ def submit_fused(result, routes, target, *, spin_cycles): # and the group refuses further plans rather than let the allocator reuse them. group._pending_release = None group._spin_bound = SimpleNamespace(cycles=lambda: _SERVING_SPIN_CYCLES) - group.broadcaster.submit_fused = lambda *args, **kwargs: SimpleNamespace(generation=99) + group.broadcaster.submit_fused = lambda *args, **kwargs: 99 stuck = _Routes(8, trace) with pytest.raises(RuntimeError, match="generation mismatch"): group.plan_schedule(stuck) From 7b39282cc41b715a7200592540c642e363e7c868 Mon Sep 17 00:00:00 2001 From: jintaop Date: Mon, 26 Oct 2026 22:59:50 -0700 Subject: [PATCH 48/59] [None][perf] Reuse the lent route view and the MegaMoE tuning op name Two more per-layer host costs of per-layer dynamic EPLB, measured in situ on a GB200 host. Outputs are unchanged: - plan_finish() sliced the scheduler's fixed physical-route buffer on every layer. Every layer of a forward plans the same token count, so the group keeps the view and reuses it while the count repeats; - a launch that may tune built the MegaMoE tuning op name, an f-string with the helper count, on every call. The name depends only on the helper count and is now built once per count. Signed-off-by: jintaop --- .../moe/custom_ops/cute_dsl_megamoe_custom_op.py | 6 +++++- .../moe/fused_moe/per_layer_dynamic_eplb/planner.py | 13 ++++++++++++- .../_torch/moe/test_dynamic_eplb_contracts.py | 13 +++++++++---- tests/unittest/_torch/moe/test_moe_backend.py | 2 ++ 4 files changed, 28 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py b/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py index ddc8cd0e371c..faac62e78056 100644 --- a/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py +++ b/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py @@ -251,8 +251,12 @@ def _token_back_ready_granularity_choices(*, load_balance: bool = False) -> tupl return _TOKEN_BACK_READY_GRANULARITIES if load_balance else ("expert",) +@functools.cache def _megamoe_tuning_op_name(helper_expert_count: int) -> str: - """Separate timing caches without changing the registered op or kernel ABI.""" + """Separate timing caches without changing the registered op or kernel ABI. + + Every launch that may tune asks for it, so each name is built once. + """ base = "trtllm::cute_dsl_megamoe_nvfp4_blackwell" load_balance = helper_expert_count > 0 alpha = _autotune_pl_alpha(helper_expert_count) diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py index af66b06e05fa..00c429880bbf 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py @@ -168,6 +168,10 @@ def __init__( self._route_wait_generation = 0 # Routes of the latest launch, held until MAIN waits for its route epoch. self._routes_in_flight: Optional[torch.Tensor] = None + # The lent physical routes view and its row count: every layer of a + # forward plans the same number of tokens, so one view serves them. + self._lent_rows = -1 + self._lent_routes: Optional[torch.Tensor] = None self._execution_stream_handle: Optional[int] = None self._closed = False self.scheduler = None @@ -286,6 +290,8 @@ def close(self) -> None: self._plan_part = None self._pending_release = None self._routes_in_flight = None + self._lent_rows = -1 + self._lent_routes = None self.lease = None self.broadcaster = None self.scheduler = None @@ -387,7 +393,12 @@ def plan_finish( outputs, num_tokens = part if not defer_wait: self.wait_for_routes() - physical = outputs.physical_slot_ids[:num_tokens] + physical = self._lent_routes + if num_tokens != self._lent_rows: + # A view of the scheduler's fixed output buffer; slicing it is + # the cost, so reuse it while the token count repeats. + physical = self._lent_routes = outputs.physical_slot_ids[:num_tokens] + self._lent_rows = num_tokens self.plan_calls += 1 self._plan_part = None return physical, self._generation diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index 49afd1b25e26..cfdd5424cffb 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -1264,6 +1264,8 @@ def test_fused_submit_orders_generations_and_defers_the_route_wait(auxiliary_sms group._plan_part = None group._pending_release = None group._routes_in_flight = None + group._lent_rows = -1 + group._lent_routes = None group._closed = False group._execution_stream_handle = None group.device = 0 @@ -1273,8 +1275,8 @@ def test_fused_submit_orders_generations_and_defers_the_route_wait(auxiliary_sms group.copy_stream = group.scheduler_stream = copy group._stream_handle = 19 group._route_epoch_ptr = 77 - # Generation 1 launches during warmup, generation 2 while serving. - bounds = iter((_WARMUP_SPIN_CYCLES, _SERVING_SPIN_CYCLES)) + # Generation 1 launches during warmup, the later ones while serving. + bounds = iter((_WARMUP_SPIN_CYCLES, _SERVING_SPIN_CYCLES, _SERVING_SPIN_CYCLES)) group._spin_bound = SimpleNamespace(cycles=lambda: next(bounds)) def driver_call(*entry): @@ -1313,7 +1315,8 @@ def submit_fused(result, routes, target, *, spin_cycles): ) group.lease = lease_cls(group) - for generation, rows in enumerate((16, 0), 1): + lent = {} + for generation, rows in enumerate((16, 16, 0), 1): trace.clear() routes = _Routes(rows, trace) target = generation * auxiliary_sms @@ -1322,6 +1325,8 @@ def submit_fused(result, routes, target, *, spin_cycles): physical, actual_generation = group.plan_finish(part, defer_wait=True) assert physical.shape == (rows, 6) assert actual_generation == generation + # The view of the fixed output buffer is reused while the token count repeats. + assert lent.setdefault(rows, physical) is physical assert route_wait not in trace # The group, not an allocator stream event, keeps the routes alive # until MAIN waits for every HALO-Q read of them. @@ -1340,7 +1345,7 @@ def submit_fused(result, routes, target, *, spin_cycles): assert trace.index(("lease", generation - 1)) < trace.index( ("FUSED", rows, target, generation) ) - assert launch_bounds == [_WARMUP_SPIN_CYCLES, _SERVING_SPIN_CYCLES] + assert launch_bounds == [_WARMUP_SPIN_CYCLES, _SERVING_SPIN_CYCLES, _SERVING_SPIN_CYCLES] # The bound-stream fast path still rejects another stream or device. main.cuda_stream = 45 diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index 59818abdf51d..b239ebad4aa6 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -2666,6 +2666,8 @@ def test_rebalance_discard_plan_orders_release_and_allows_next_reuse(): group = RebalanceSlotSchedulerGroup.__new__(RebalanceSlotSchedulerGroup) group._check_owner = MagicMock() + group._lent_rows = -1 + group._lent_routes = None group._plan_part = part group._generation = 1 group._finished_generation = 0 From 33fa7a410cf582c6cd0119e404c96d2533851bee Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 27 Oct 2026 07:56:05 -0700 Subject: [PATCH 49/59] [None][perf] Cut per-request host overhead on rank 0 In short-prompt prefill, rank 0 handles about a thousand requests per iteration, and its host path can pace every rank. Outputs, messages and per-module log levels are unchanged: - SamplingParams builds the SamplingConfig and OutputConfig field sets once instead of scanning dir() of the bindings on every request, and pickles only its field values in field order, which the rank-0 worker restores about 3x faster on the thread it shares with the executor loop; - two per-iteration executor debug messages are formatted only when debug logging is enabled; - Logger.log() and is_debug_enabled() skip the caller-module stack walk when no per-module level is configured and the global level already filters the message; - with TRTLLM_PROXY_BATCH_REQUESTS=1 (default off), the executor proxy pickles each request on the submitting thread and one background thread sends each burst to the rank-0 worker as one HMAC-signed message, which the worker unpacks; the same thread then sends the shutdown sentinel after the last request, and it is not joined during interpreter teardown; - with TRTLLM_BENCH_DISABLE_GC=1 (default off), trtllm-bench throughput and latency collect once after warmup, freeze the surviving objects and disable the cyclic GC for the measured run, like TRTLLM_WORKER_DISABLE_GC and TRTLLM_SERVER_DISABLE_GC do in the workers and the server. These changes do not depend on the MoE load balancer. Signed-off-by: jintaop --- .../developer-guide/perf-benchmarking.md | 21 ++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 13 +- tensorrt_llm/bench/benchmark/low_latency.py | 5 +- tensorrt_llm/bench/benchmark/throughput.py | 5 +- tensorrt_llm/bench/benchmark/utils/general.py | 31 +- tensorrt_llm/executor/proxy.py | 180 +++++++++++- tensorrt_llm/executor/worker.py | 18 +- tensorrt_llm/logger.py | 6 + tensorrt_llm/sampling_params.py | 51 +++- .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../executor/test_proxy_request_sender.py | 267 ++++++++++++++++++ tests/unittest/llmapi/test_sampling_params.py | 18 ++ tests/unittest/utils/test_logger.py | 21 ++ 13 files changed, 612 insertions(+), 25 deletions(-) create mode 100644 tests/unittest/executor/test_proxy_request_sender.py diff --git a/docs/source/developer-guide/perf-benchmarking.md b/docs/source/developer-guide/perf-benchmarking.md index 0f1db502d914..9738e4c0ceb0 100644 --- a/docs/source/developer-guide/perf-benchmarking.md +++ b/docs/source/developer-guide/perf-benchmarking.md @@ -37,6 +37,7 @@ The following guidance will mostly focus on benchmarks using `trtllm-bench` CLI. - [Supported Quantization Modes](#supported-quantization-modes) - [Preparing a Dataset](#preparing-a-dataset) - [Running with the PyTorch Workflow](#running-with-the-pytorch-workflow) + - [Keeping Host Stalls out of the Measured Run](#keeping-host-stalls-out-of-the-measured-run) - [Benchmarking with LoRA Adapters in PyTorch workflow](#benchmarking-with-lora-adapters-in-pytorch-workflow) - [Running multi-modal models in the PyTorch Workflow](#running-multi-modal-models-in-the-pytorch-workflow) - [Quantization in the PyTorch Flow](#quantization-in-the-pytorch-flow) @@ -251,6 +252,26 @@ trtllm-bench --model meta-llama/Llama-3.1-8B \ --backend pytorch ``` +#### Keeping Host Stalls out of the Measured Run + +With many short requests, for example context-only runs at high concurrency, +rank 0 can be bound by host work, and a pause there makes every other rank wait +at the next collective. With `TRTLLM_BENCH_DISABLE_GC=1`, the `throughput` and +`latency` commands collect once after warmup, freeze the surviving objects, and +disable Python's cyclic garbage collector for the measured run, so a full +collection over the in-flight requests cannot pause the benchmark process. +`TRTLLM_WORKER_DISABLE_GC=1` does the same in the worker processes. Memory held +in reference cycles is then not freed until the run ends, so set the variables +identically in every run you compare. + +With `TRTLLM_PROXY_BATCH_REQUESTS=1`, the executor proxy pickles each request on +the submitting thread and a background thread sends each burst to the rank-0 +worker as one message, so the worker receives and authenticates one message per +burst instead of one per request. The sender thread must take the GIL first, so +while the submitting thread keeps running Python, as a busy `trtllm-serve` event +loop does, a lone request can wait up to the interpreter's switch interval (5 ms +by default) before it is sent. + #### Benchmarking with LoRA Adapters in PyTorch workflow The PyTorch workflow supports benchmarking with LoRA (Low-Rank Adaptation) adapters. This requires preparing a dataset with LoRA metadata and configuring the LoRA settings. diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index bffe50156091..28b3792ec224 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -10012,8 +10012,10 @@ def _enqueue_responses(self, responses: Iterable[Tuple[int, LlmResponse]]): f"_enqueue_responses.") gather_responses.extend(resp) responses = gather_responses - logger.debug( - f'after gather, rank = {self.dist.rank}, responses = {responses}') + if logger.is_debug_enabled(): + logger.debug( + f'after gather, rank = {self.dist.rank}, responses = {responses}' + ) if self.dist.rank == 0 or self.gather_all_responses: with self.response_cv: @@ -10121,9 +10123,10 @@ def _handle_responses(self, emit_first_iter: bool = True): requests_finished_by_transfer = [] new_active_requests = [] timed_out_requests = [] - logger.debug( - f'------before _handle_responses, rank = {self.dist.rank}, output = {self.active_requests}' - ) + if logger.is_debug_enabled(): + logger.debug( + f'------before _handle_responses, rank = {self.dist.rank}, output = {self.active_requests}' + ) batch_token_time = self.perf_manager.get_timestamp() gpu_times_cache = {} diff --git a/tensorrt_llm/bench/benchmark/low_latency.py b/tensorrt_llm/bench/benchmark/low_latency.py index 25fa2ee9e828..96a142ad2932 100644 --- a/tensorrt_llm/bench/benchmark/low_latency.py +++ b/tensorrt_llm/bench/benchmark/low_latency.py @@ -27,7 +27,8 @@ generate_json_report, get_general_cli_options, get_llm) from tensorrt_llm.bench.benchmark.utils.asynchronous import async_benchmark -from tensorrt_llm.bench.benchmark.utils.general import generate_warmup_dataset +from tensorrt_llm.bench.benchmark.utils.general import (generate_warmup_dataset, + measurement_gc) from tensorrt_llm.bench.dataclasses.configuration import RuntimeConfig from tensorrt_llm.bench.dataclasses.general import BenchmarkEnvironment from tensorrt_llm.bench.dataclasses.reporting import ReportUtility @@ -339,7 +340,7 @@ def latency_command( logger.info("Warmup done.") iteration_writer = options.iteration_writer - with iteration_writer.capture(): + with iteration_writer.capture(), measurement_gc(): statistics = asyncio.run( async_benchmark(llm, sampling_params, diff --git a/tensorrt_llm/bench/benchmark/throughput.py b/tensorrt_llm/bench/benchmark/throughput.py index a28d6e00a2b2..22a8f55fb53e 100755 --- a/tensorrt_llm/bench/benchmark/throughput.py +++ b/tensorrt_llm/bench/benchmark/throughput.py @@ -36,7 +36,8 @@ get_settings) # isort: on from tensorrt_llm.bench.benchmark.utils.general import ( - generate_warmup_dataset, update_sampler_args_with_extra_options) + generate_warmup_dataset, measurement_gc, + update_sampler_args_with_extra_options) from tensorrt_llm.bench.dataclasses.configuration import RuntimeConfig from tensorrt_llm.bench.dataclasses.general import BenchmarkEnvironment from tensorrt_llm.bench.dataclasses.reporting import ReportUtility @@ -494,7 +495,7 @@ def throughput_command( logger.info("Warmup done.") iteration_writer = options.iteration_writer - with iteration_writer.capture(): + with iteration_writer.capture(), measurement_gc(): statistics = asyncio.run( async_benchmark(llm, sampling_params, diff --git a/tensorrt_llm/bench/benchmark/utils/general.py b/tensorrt_llm/bench/benchmark/utils/general.py index cf567b89a68b..e98f90337c48 100755 --- a/tensorrt_llm/bench/benchmark/utils/general.py +++ b/tensorrt_llm/bench/benchmark/utils/general.py @@ -1,8 +1,11 @@ from __future__ import annotations +import gc +import os +from contextlib import contextmanager from pathlib import Path from random import choices, shuffle -from typing import Dict, List, Union +from typing import Dict, Iterator, List, Union import yaml @@ -168,6 +171,32 @@ def get_settings(params: dict, dataset_metadata: DatasetMetadata, model: str, } +@contextmanager +def measurement_gc() -> Iterator[None]: + """Keep Python's cyclic GC out of the measured run if TRTLLM_BENCH_DISABLE_GC=1. + + The benchmark process submits every request and handles every response. A + full collection over its in-flight and finished requests can stop it long + enough for the executors to run out of requests. This collects once, + freezes the surviving objects and disables the GC until the measured run + ends, as TRTLLM_SERVER_DISABLE_GC and TRTLLM_WORKER_DISABLE_GC disable it + in the server and the workers. + """ + if os.getenv("TRTLLM_BENCH_DISABLE_GC", "0") != "1": + yield + return + gc.collect() + gc.freeze() + gc.disable() + logger.info("Python GC disabled for the measured run " + "(TRTLLM_BENCH_DISABLE_GC=1).") + try: + yield + finally: + gc.enable() + gc.unfreeze() + + def generate_warmup_dataset(requests, steps) -> List[InferenceRequest]: """Warm up the benchmarker.""" warm_up_dataset = choices(requests, k=steps) diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index 68692a0adddc..3e0f3f919ecc 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/tensorrt_llm/executor/proxy.py @@ -13,15 +13,18 @@ # See the License for the specific language governing permissions and # limitations under the License. import atexit +import collections import concurrent.futures import json import os +import pickle # nosec B403 import shutil +import sys import tempfile import threading import time import weakref -from queue import Empty +from queue import Empty, Queue from typing import Dict, List, Optional, Union import torch @@ -100,9 +103,132 @@ def _check_collective_rpc_guard( "this shim only reaches rank-0.") +class _RequestSender: + """Batch a proxy's requests to the rank-0 worker on one background thread. + + Used only with ``TRTLLM_PROXY_BATCH_REQUESTS=1``. Callers pickle each + request themselves and ``send`` the bytes, so a request is captured as it + is when it is submitted, and one that cannot be pickled fails in its + caller. The sender thread sends every payload queued since its previous + send as one list in one HMAC-signed message, so a burst of requests costs + the rank-0 worker one receive and one HMAC check instead of one per + request; the worker then unpickles each payload. The thread starts on + the first ``send``. From then on it is the only thread that writes to the + request queue's socket, which is not thread-safe, and it also sends the + worker's shutdown sentinel. + + Trade-off: the sender thread must take the GIL before it can send. While + the submitting thread keeps running Python, for example a busy + trtllm-serve event loop, a lone request can wait up to the interpreter's + switch interval (5 ms by default) before it is sent. + """ + + MAX_BATCH_SIZE = 1024 + + def __init__(self, request_queue: IpcQueue, error_queue: Queue): + self._request_queue = request_queue + self._error_queue = error_queue + self._pending = collections.deque() + self._cond = threading.Condition() + self._closed = False + self._send_sentinel = False + # Set by the thread once it has drained the queue and will stop. + self._stopping = False + self._thread: Optional[threading.Thread] = None + # Bound here because ``close`` can run from a finalizer during + # interpreter teardown, when module globals may already be None. + self._is_finalizing = sys.is_finalizing + + def send(self, payload: bytes) -> None: + """Queue a pickled request for the worker. + + Payloads sent after ``close`` are dropped: the proxy is shutting down. + """ + with self._cond: + if self._closed: + return + if self._thread is None: + self._thread = threading.Thread(target=self._run, + daemon=True, + name="proxy_request_sender") + self._thread.start() + self._pending.append(payload) + self._cond.notify() + + def close(self, + timeout: Optional[float] = None, + *, + send_sentinel: bool = False) -> bool: + """Send everything queued, then stop the sender thread. + + With ``send_sentinel``, the thread also sends the worker's shutdown + sentinel after the last payload, so the request queue keeps a single + writer. Calls after the thread has stopped return at once. During + interpreter finalization this returns without doing anything: the + daemon thread can no longer run, the ZMQ context may be gone, and + joining a thread there fails. + + Returns: + Whether the caller still has to send the sentinel: true only if + ``send_sentinel`` is set and the thread did not send it, because it + never started or had already stopped. + """ + if self._is_finalizing(): + return False + with self._cond: + self._closed = True + thread = self._thread + caller_sends_sentinel = send_sentinel and (thread is None + or self._stopping) + self._send_sentinel = self._send_sentinel or send_sentinel + self._cond.notify() + if thread is not None: + thread.join(timeout) + if not thread.is_alive(): + self._thread = None + return caller_sends_sentinel + + def _run(self) -> None: + while True: + with self._cond: + while not self._pending and not self._closed: + self._cond.wait() + if not self._pending: + self._stopping = True + send_sentinel = self._send_sentinel + break + batch = [ + self._pending.popleft() + for _ in range(min(len(self._pending), self.MAX_BATCH_SIZE)) + ] + try: + self._request_queue.put(batch) + except Exception as e: + # The payloads are already pickled, so this is a transport + # error, which the proxy treats as fatal. + self._error_queue.put_nowait(e) + if send_sentinel: + # A bare None, not a batch: the worker stops at a top-level None. + try: + self._request_queue.put_noblock(None, retry=4) + except Exception as e: + self._error_queue.put_nowait(e) + + +def _create_request_sender(request_queue: IpcQueue, + error_queue: Queue) -> Optional[_RequestSender]: + """Return a batching request sender if TRTLLM_PROXY_BATCH_REQUESTS=1.""" + if os.getenv("TRTLLM_PROXY_BATCH_REQUESTS", "0") != "1": + return None + return _RequestSender(request_queue, error_queue) + + class GenerationExecutorProxy(GenerationExecutor): READY_SIGNAL = b"READY" WORKER_PROCESS_IDENTITIES_SIGNAL = b"WORKER_PROCESS_IDENTITIES" + # Set when TRTLLM_PROXY_BATCH_REQUESTS=1; requests are put on + # ``request_queue`` directly otherwise. + _request_sender: Optional[_RequestSender] = None # Total wait for the worker futures in an orderly shutdown. Generous # enough for the teardown of a large engine. _WORKER_SHUTDOWN_TIMEOUT_S = 180.0 @@ -258,6 +384,9 @@ def __init__( # ``abort_request()`` on user-caller threads — that pre-existing # cross-thread access is out of scope for this change. The fix # here narrowly targets the new profile-control path. + # With TRTLLM_PROXY_BATCH_REQUESTS=1, the ``_request_sender`` thread + # writes the queue instead once it has started, shutdown sentinel + # included. self._profile_control_executor = concurrent.futures.ThreadPoolExecutor( max_workers=1, thread_name_prefix="proxy_profile_control") @@ -487,6 +616,8 @@ def _setup_queues(self) -> WorkerCommIpcAddrs: is_server=False, socket_type=zmq.PUSH, name="proxy_request_queue") + self._request_sender = _create_request_sender( + self.request_queue, self._error_queue) self.result_queue = FusedIpcQueue(frontend_result_addrs[0], is_server=True, fuse_message=False, @@ -496,6 +627,8 @@ def _setup_queues(self) -> WorkerCommIpcAddrs: request_addr = None self.request_queue = IpcQueue(is_server=True, name="proxy_request_queue") + self._request_sender = _create_request_sender( + self.request_queue, self._error_queue) # TODO[chunweiy]: Unify IpcQueue and FusedIpcQueue # Use PULL mode when enable_postprocess_parallel as there are # multiple senders from multiple processes. @@ -584,7 +717,7 @@ def abort_request(self, request_id: int) -> None: # NOTE, it just sends a cancelling request to the request queue, but it # may take a while for the request to be cancelled in the worker and # send back a finished result. - self.request_queue.put(CancellingRequest(request_id)) + self._send_request(CancellingRequest(request_id)) # Per-call timeouts for the synchronous worker round-trip. The stop # path can take up to 30s because PyExecutor.stop_profile itself @@ -662,10 +795,12 @@ def _profile_control_call(self, kind: str, request, timeout: float) -> None: Runs on ``_profile_control_executor`` so all profile-related ZMQ socket ops (``request_queue.put`` and ``profile_ack_queue.get`` - via ``_wait_profile_ack``) happen on a single owning thread. + via ``_wait_profile_ack``) happen on a single owning thread. With + TRTLLM_PROXY_BATCH_REQUESTS=1 the request is put by + ``_request_sender`` instead. """ request.ack_addr = self._ensure_profile_ack_queue() - self.request_queue.put(request) + self._send_request(request) self._wait_profile_ack(kind, timeout) def _submit_profile_control(self, kind: str, request, @@ -1047,7 +1182,35 @@ def pre_shutdown(self): # empty list and therefore covered case (a) by accident. Switching # to bare ``any(...)`` to better handle partial crashes regressed # case (a); keep both behaviours explicit. - if not self.mpi_futures or any(not f.done() for f in self.mpi_futures): + send_sentinel = not self.mpi_futures or any(not f.done() + for f in self.mpi_futures) + self._close_request_sender(send_sentinel=send_sentinel) + + def _send_request(self, request) -> None: + """Send a request to the rank-0 worker. + + With TRTLLM_PROXY_BATCH_REQUESTS=1 the request is pickled here, on the + caller's thread, and ``_request_sender`` batches the bytes; otherwise + it is put on ``request_queue`` directly. Either way the worker receives + the request as it is now, and a request that cannot be pickled raises + here. + """ + if self._request_sender is None: + self.request_queue.put(request) + else: + self._request_sender.send(pickle.dumps(request)) + + def _close_request_sender(self, send_sentinel: bool = False) -> None: + """Stop the request sender, if any, then send the shutdown sentinel if asked. + + The sentinel follows every queued request. It is sent by the thread + that has been writing ``request_queue``: the sender thread once it has + started, the caller otherwise. + """ + if self._request_sender is not None: + send_sentinel = self._request_sender.close( + timeout=5.0, send_sentinel=send_sentinel) + if send_sentinel: self.request_queue.put_noblock(None, retry=4) def _get_next_client_id(self) -> int: @@ -1072,6 +1235,7 @@ def _cleanup_multi_frontend_ipc_dir(self): def shutdown(self): if not self.workers_started: + self._close_request_sender() self._cleanup_multi_frontend_ipc_dir() return @@ -1210,7 +1374,7 @@ def submit(self, request: GenerationRequest) -> GenerationResult: raise EngineDeadError(self._fatal_error) with nvtx_range_debug("request_queue.put"): - self.request_queue.put(request) + self._send_request(request) self._handle_background_error() @@ -1530,6 +1694,8 @@ def __init__( is_server=False, socket_type=zmq.PUSH, name=f"frontend_{frontend_id}_request_queue") + self._request_sender = _create_request_sender(self.request_queue, + self._error_queue) self.result_queue = FusedIpcQueue( (attach_info["result_addrs"][frontend_id], hmac_key), is_server=True, @@ -1579,6 +1745,8 @@ def pre_shutdown(self): def shutdown(self): self.pre_shutdown() + # Send the cancels queued by pre_shutdown before the lane goes idle. + self._close_request_sender() if self.rpc_client is not None: self.rpc_client.close() self.rpc_client = None diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py index 9c25aaf88f3f..602438cf13a7 100644 --- a/tensorrt_llm/executor/worker.py +++ b/tensorrt_llm/executor/worker.py @@ -15,6 +15,7 @@ import gc import os +import pickle # nosec B403 import sys import threading import time @@ -358,6 +359,21 @@ def block_subordinates(self): self.engine.wait_shutdown() +def _incoming_requests(request_queue: IpcQueue): + """Yield the proxy's requests until its shutdown sentinel arrives. + + With TRTLLM_PROXY_BATCH_REQUESTS=1 the proxy sends a burst of requests as + one list of pickled requests. ``request_queue.get()`` has already + authenticated that message, and with it every payload inside. + """ + while (msg := request_queue.get()) is not None: + if isinstance(msg, list): + for payload in msg: + yield pickle.loads(payload) # nosec B301 + else: + yield msg + + @print_traceback_on_error def worker_main( worker_queues: WorkerCommIpcAddrs, @@ -630,7 +646,7 @@ def notify_proxy_threads_to_quit(): worker.engine.set_resource_governor_queue( resource_governor_queue) - while (req := request_queue.get()) is not None: + for req in _incoming_requests(request_queue): if isinstance(req, CancellingRequest): worker.abort_request(req.id) elif isinstance(req, StartProfileRequest): diff --git a/tensorrt_llm/logger.py b/tensorrt_llm/logger.py index 74253479a9b8..35fa5b6d077e 100644 --- a/tensorrt_llm/logger.py +++ b/tensorrt_llm/logger.py @@ -296,6 +296,10 @@ def _percent_format(cls, msg) -> Optional[str]: return None def log(self, severity, *msg): + # Without per-module levels the caller's module only matters for an emitted message; + # a filtered one returns before the stack walk. + if not self._module_levels and not self.is_severity_enabled(severity): + return module = _get_caller_module() if not self.is_severity_enabled(severity, module): return @@ -364,6 +368,8 @@ def is_debug_enabled(self) -> bool: ``TLLM_LOG_LEVEL_BY_MODULE`` override can enable or disable the caller's module independently of the global level. """ + if not self._module_levels: + return self.is_severity_enabled(self.VERBOSE) return self.is_severity_enabled(self.VERBOSE, _get_caller_module()) @property diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index c133f71946f4..49b28ebc679a 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -12,11 +12,13 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import functools import json +import operator import os from abc import ABC, abstractmethod from collections.abc import Iterable, Mapping -from dataclasses import dataclass, field, fields +from dataclasses import dataclass, field, fields, make_dataclass from typing import Any, List, NamedTuple, Optional, Tuple, Union import torch @@ -375,6 +377,12 @@ def __post_init__(self): self._validate() + def __reduce_ex__(self, protocol): + # Pickle the field values in field order; subclasses may add fields. + if type(self) is not SamplingParams: + return object.__reduce_ex__(self, protocol) + return _restore_sampling_params, (_get_sampling_params_state(self),) + def _validate(self): """Verify the sampling parameters. @@ -669,7 +677,7 @@ def _get_sampling_config(self) -> tllme.SamplingConfig: # | Sampling | use_beam_search | beam_width == 1 | # | Sampling | n | num_return_sequences | # | Sampling | best_of | no corresponding param | - fields = {f for f in dir(tllme.SamplingConfig) if not f.startswith("__")} + fields = _sampling_config_fields() unmatched_params = [ "num_return_sequences", "beam_width", @@ -688,12 +696,7 @@ def _get_sampling_config(self) -> tllme.SamplingConfig: return tllme.SamplingConfig(**llmapi_to_rt_param_map) def _get_output_config(self) -> tllme.OutputConfig: - sampling_param_fields = set(dir(SamplingParams)) - fields = [ - f - for f in dir(tllme.OutputConfig) - if not f.startswith("__") and f in sampling_param_fields - ] + fields = _output_config_fields() config_kwargs = {f: getattr(self, f) for f in fields} @@ -751,3 +754,35 @@ def _get_guided_decoding_params(self) -> tllme.GuidedDecodingParams: ) else: return None + + +_SAMPLING_PARAMS_FIELDS = tuple(f.name for f in fields(SamplingParams)) +_get_sampling_params_state = operator.attrgetter(*_SAMPLING_PARAMS_FIELDS) +# Only its generated ``__init__`` is used: one store per field, no ``__post_init__``. +_set_sampling_params_state = make_dataclass( + "_SamplingParamsState", _SAMPLING_PARAMS_FIELDS, repr=False, eq=False +).__init__ + + +def _restore_sampling_params(state: tuple) -> "SamplingParams": + sampling_params = object.__new__(SamplingParams) + _set_sampling_params_state(sampling_params, *state) + return sampling_params + + +@functools.cache +def _sampling_config_fields() -> frozenset[str]: + """Return the ``tllme.SamplingConfig`` field names. + + They are fixed by the bindings, so the ``dir()`` scan runs once instead of once per request. + """ + return frozenset(f for f in dir(tllme.SamplingConfig) if not f.startswith("__")) + + +@functools.cache +def _output_config_fields() -> tuple[str, ...]: + """Return the ``tllme.OutputConfig`` field names that ``SamplingParams`` also defines.""" + sampling_param_fields = set(dir(SamplingParams)) + return tuple( + f for f in dir(tllme.OutputConfig) if not f.startswith("__") and f in sampling_param_fields + ) diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index a02f527b12d2..db3f1ffc42ea 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -203,6 +203,7 @@ l0_cpu: - unittest/executor/test_ipc.py - unittest/executor/test_rpc.py - unittest/executor/test_multi_frontend_routing.py + - unittest/executor/test_proxy_request_sender.py - unittest/executor/test_event_loop_error_broadcast.py - unittest/executor/test_forward_error_survival.py - unittest/executor/test_rpc_worker_mixin.py diff --git a/tests/unittest/executor/test_proxy_request_sender.py b/tests/unittest/executor/test_proxy_request_sender.py new file mode 100644 index 000000000000..7dbeb489063d --- /dev/null +++ b/tests/unittest/executor/test_proxy_request_sender.py @@ -0,0 +1,267 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for request batching between the proxy and the rank-0 worker. + +``TRTLLM_PROXY_BATCH_REQUESTS=1`` makes the proxy pickle each request on the +submitting thread and send bursts of payloads as one list; the worker unpickles +each payload. +""" + +import pickle +import queue +import threading +from types import SimpleNamespace + +import pytest + +from tensorrt_llm.executor.proxy import ( + GenerationExecutorProxy, + _create_request_sender, + _RequestSender, +) +from tensorrt_llm.executor.worker import _incoming_requests + +pytestmark = pytest.mark.cpu_only + + +class _RecordingQueue: + """Request-queue stand-in that records each message; the first put can block. + + Like ``ZeroMqQueue`` with ``TLLM_LLMAPI_ZMQ_DEBUG=1``, it rejects a write from + any thread other than the first one that wrote. + """ + + def __init__(self, block_first_put: bool = False): + self.messages = [] + self.writer = None + self.first_put_started = threading.Event() + self.release_first_put = threading.Event() + if not block_first_put: + self.release_first_put.set() + + def _check_writer(self): + if self.writer is None: + self.writer = threading.get_ident() + elif self.writer != threading.get_ident(): + raise RuntimeError("request queue written from a second thread") + + def put(self, msg): + self._check_writer() + if not self.messages: + self.first_put_started.set() + self.release_first_put.wait(timeout=10) + self.messages.append(msg) + + def put_noblock(self, msg, *, retry=1, wait_time=0.001): + self._check_writer() + self.messages.append(msg) + + +def _payloads(*items): + return [pickle.dumps(item) for item in items] + + +def _bare_proxy(): + proxy = GenerationExecutorProxy.__new__(GenerationExecutorProxy) + # Keeps the __del__ -> shutdown() path a no-op at GC time. + proxy.workers_started = False + return proxy + + +def _pre_shutdown(proxy): + """Run ``pre_shutdown`` on a bare proxy with live workers, on a new thread.""" + proxy.workers_started = True + proxy.doing_shutdown = False + proxy._worker_process_monitor = SimpleNamespace(close=lambda: None) + proxy._results = {} + proxy.mpi_futures = [] + shutdown = threading.Thread(target=proxy.pre_shutdown) + shutdown.start() + return shutdown + + +def test_batching_is_opt_in(monkeypatch): + monkeypatch.delenv("TRTLLM_PROXY_BATCH_REQUESTS", raising=False) + assert _create_request_sender(_RecordingQueue(), queue.Queue()) is None + monkeypatch.setenv("TRTLLM_PROXY_BATCH_REQUESTS", "1") + assert isinstance(_create_request_sender(_RecordingQueue(), queue.Queue()), _RequestSender) + + +def test_lone_request_is_sent_as_one_payload_list(): + request_queue = _RecordingQueue() + sender = _RequestSender(request_queue, queue.Queue()) + sender.send(pickle.dumps("a")) + sender.close(timeout=10) + assert request_queue.messages == [_payloads("a")] + + +def test_burst_is_sent_as_one_list_in_order(): + request_queue = _RecordingQueue(block_first_put=True) + sender = _RequestSender(request_queue, queue.Queue()) + sender.send(pickle.dumps(0)) + assert request_queue.first_put_started.wait(timeout=10) + # Everything queued while the first send is in flight goes out as one list. + for item in range(1, 6): + sender.send(pickle.dumps(item)) + request_queue.release_first_put.set() + sender.close(timeout=10) + assert request_queue.messages == [_payloads(0), _payloads(1, 2, 3, 4, 5)] + + +def test_close_flushes_queued_requests_and_drops_later_ones(): + request_queue = _RecordingQueue(block_first_put=True) + sender = _RequestSender(request_queue, queue.Queue()) + sender.send(pickle.dumps("first")) + assert request_queue.first_put_started.wait(timeout=10) + sender.send(pickle.dumps("queued")) + closer = threading.Thread(target=sender.close, kwargs={"timeout": 10}) + closer.start() + request_queue.release_first_put.set() + closer.join(timeout=10) + assert not closer.is_alive() + assert request_queue.messages == [_payloads("first"), _payloads("queued")] + sender.send(pickle.dumps("late")) + assert request_queue.messages == [_payloads("first"), _payloads("queued")] + + +def test_close_is_idempotent(): + request_queue = _RecordingQueue() + sender = _RequestSender(request_queue, queue.Queue()) + sender.send(pickle.dumps("a")) + sender.close(timeout=10) + assert sender._thread is None + sender.close(timeout=10) + assert request_queue.messages == [_payloads("a")] + + +def test_sender_thread_sends_the_sentinel_after_the_queued_requests(): + request_queue = _RecordingQueue(block_first_put=True) + sender = _RequestSender(request_queue, queue.Queue()) + sender.send(pickle.dumps("first")) + assert request_queue.first_put_started.wait(timeout=10) + sender.send(pickle.dumps("queued")) + result = [] + closer = threading.Thread( + target=lambda: result.append(sender.close(timeout=10, send_sentinel=True)) + ) + closer.start() + request_queue.release_first_put.set() + closer.join(timeout=10) + assert result == [False] + # A bare None after the payload lists, written by the sender thread. + assert request_queue.messages == [_payloads("first"), _payloads("queued"), None] + + +def test_caller_sends_the_sentinel_when_the_sender_thread_did_not(): + # The thread never started, so nothing has written the queue yet. + assert _RequestSender(_RecordingQueue(), queue.Queue()).close(send_sentinel=True) + # The thread stopped before the sentinel was asked for. + sender = _RequestSender(_RecordingQueue(), queue.Queue()) + sender.send(pickle.dumps("a")) + sender.close(timeout=10) + assert sender.close(timeout=10, send_sentinel=True) + + +def test_close_during_interpreter_finalization_does_not_join(): + request_queue = _RecordingQueue(block_first_put=True) + sender = _RequestSender(request_queue, queue.Queue()) + sender.send(pickle.dumps("first")) + assert request_queue.first_put_started.wait(timeout=10) + sender.send(pickle.dumps("queued")) + # Joining would wait for the blocked send; at teardown close() must return at + # once, and must not ask anyone to send the sentinel. + sender._is_finalizing = lambda: True + assert sender.close(send_sentinel=True) is False + assert sender._thread is not None + # Outside finalization the same sender still flushes and stops. + sender._is_finalizing = lambda: False + request_queue.release_first_put.set() + sender.close(timeout=10) + assert sender._thread is None + assert request_queue.messages == [_payloads("first"), _payloads("queued")] + + +def test_pre_shutdown_flushes_batched_requests_before_the_sentinel(): + request_queue = _RecordingQueue(block_first_put=True) + proxy = _bare_proxy() + proxy.request_queue = request_queue + proxy._request_sender = _RequestSender(request_queue, queue.Queue()) + proxy._send_request("first") + assert request_queue.first_put_started.wait(timeout=10) + proxy._send_request("queued") + + shutdown = _pre_shutdown(proxy) + request_queue.release_first_put.set() + shutdown.join(timeout=10) + assert not shutdown.is_alive() + proxy.workers_started = False + # The sender thread wrote everything, the sentinel included. + assert request_queue.messages == [_payloads("first"), _payloads("queued"), None] + assert request_queue.writer != shutdown.ident + + +@pytest.mark.parametrize("batching", [True, False]) +def test_pre_shutdown_sends_the_sentinel_itself_if_nothing_else_wrote(batching): + request_queue = _RecordingQueue() + proxy = _bare_proxy() + proxy.request_queue = request_queue + if batching: + proxy._request_sender = _RequestSender(request_queue, queue.Queue()) + + shutdown = _pre_shutdown(proxy) + shutdown.join(timeout=10) + assert not shutdown.is_alive() + proxy.workers_started = False + assert request_queue.messages == [None] + assert request_queue.writer == shutdown.ident + + +def test_transport_error_reaches_error_queue(): + class _FailingQueue: + def put(self, msg): + raise RuntimeError("socket closed") + + error_queue = queue.Queue() + sender = _RequestSender(_FailingQueue(), error_queue) + sender.send(pickle.dumps("a")) + sender.close(timeout=10) + error = error_queue.get_nowait() + assert isinstance(error, RuntimeError) + assert str(error) == "socket closed" + + +def test_batched_request_is_captured_when_sent(): + proxy = _bare_proxy() + sent = [] + proxy._request_sender = SimpleNamespace(send=sent.append) + + request = {"max_tokens": 1} + proxy._send_request(request) + # A change after submission does not reach the worker. + request["max_tokens"] = 2 + assert pickle.loads(sent[0]) == {"max_tokens": 1} + + # A request that cannot be pickled fails in its caller, and nothing is queued. + with pytest.raises(TypeError): + proxy._send_request({"lock": threading.Lock()}) + assert len(sent) == 1 + + +def test_unbatched_request_is_put_directly(): + proxy = _bare_proxy() + request_queue = _RecordingQueue() + proxy.request_queue = request_queue + proxy._send_request("a") + assert request_queue.messages == ["a"] + + +def test_worker_unpacks_batched_requests_in_order(): + class _ListQueue: + def __init__(self, messages): + self._messages = list(messages) + + def get(self): + return self._messages.pop(0) + + request_queue = _ListQueue(["a", _payloads("b", "c"), "d", None, "after sentinel"]) + assert list(_incoming_requests(request_queue)) == ["a", "b", "c", "d"] diff --git a/tests/unittest/llmapi/test_sampling_params.py b/tests/unittest/llmapi/test_sampling_params.py index 3c42fc097cf5..4a10387338d2 100644 --- a/tests/unittest/llmapi/test_sampling_params.py +++ b/tests/unittest/llmapi/test_sampling_params.py @@ -13,7 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. import asyncio +import copy import json +import pickle +from dataclasses import fields from typing import Any import pytest @@ -301,6 +304,21 @@ def test_sampling_params_accepts_in_range_values(field, value): assert getattr(SamplingParams(**{field: value}), field) == value +@pytest.mark.parametrize( + "round_trip", + [lambda params: pickle.loads(pickle.dumps(params)), copy.copy, copy.deepcopy], +) +def test_sampling_params_copy_round_trip_keeps_every_field(round_trip): + # A distinct value in every field, so a field the pickling misses fails. + params = SamplingParams() + for f in fields(SamplingParams): + object.__setattr__(params, f.name, f"value of {f.name}") + restored = round_trip(params) + assert type(restored) is SamplingParams + for f in fields(SamplingParams): + assert getattr(restored, f.name) == getattr(params, f.name), f.name + + @pytest.mark.parametrize("value", [None, -1]) def test_thinking_token_budget_unlimited_values(value): assert SamplingParams(thinking_token_budget=value).thinking_token_budget is None diff --git a/tests/unittest/utils/test_logger.py b/tests/unittest/utils/test_logger.py index 27a209e9d801..7d65a1182432 100644 --- a/tests/unittest/utils/test_logger.py +++ b/tests/unittest/utils/test_logger.py @@ -14,6 +14,7 @@ # limitations under the License. """Unit tests for the modular logger (auto module detection, per-module filtering).""" +import importlib import logging from io import StringIO @@ -239,6 +240,26 @@ def test_unmatched_module_uses_global(self): # Cleanup singleton._module_levels = {} + def test_filtered_message_skips_caller_lookup(self, monkeypatch): + # The package re-exports the Logger singleton as ``tensorrt_llm.logger``, + # shadowing the submodule, so fetch the module object explicitly. + logger_module = importlib.import_module("tensorrt_llm.logger") + + singleton = Logger() + monkeypatch.setattr(singleton, "_min_severity", "warning") + monkeypatch.setattr(singleton, "_module_levels", {}) + + def fail(): + raise AssertionError("caller module looked up for a filtered message") + + monkeypatch.setattr(logger_module, "_get_caller_module", fail) + singleton.info("filtered") + assert singleton.is_debug_enabled() is False + # A per-module override still needs the caller's module. + monkeypatch.setattr(singleton, "_module_levels", {"runtime": 10}) + with pytest.raises(AssertionError, match="caller module"): + singleton.info("module filtered") + class TestIsDebugEnabled: """Tests for is_debug_enabled, the guard for expensive debug arguments. From f33658e76e59a3a9b03db148a6f8a112b2898f1c Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 27 Oct 2026 07:58:38 -0700 Subject: [PATCH 50/59] [None][test] Follow main's executor changes in the dynamic EPLB contract tests main #20996 renamed PyTorchModelEngine to ModelEngine, and the source-contract tests extract the class by that name. main #19369 stops the sampler's async worker before the CUDA graphs are released and reads enable_attention_dp and _kv_connector_failed during shutdown, so the shutdown-retry test now expects the sampler stop in the runtime phase and sets those attributes. Signed-off-by: jintaop --- .../_torch/moe/test_dynamic_eplb_contracts.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index cfdd5424cffb..0111b2a3f9e7 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -601,7 +601,7 @@ def test_model_engine_warmup_uses_load_balancer_public_api(): events = [] engine_cls = _extract_class( _TORCH_ROOT / "pyexecutor/model_engine.py", - "PyTorchModelEngine", + "ModelEngine", {"_set_warmup_effects", "warmup_scope"}, { "contextmanager": contextmanager, @@ -633,7 +633,7 @@ def test_model_engine_cleanup_shuts_manager_after_graphs_while_model_is_alive(): events = [] engine_cls = _extract_class( _TORCH_ROOT / "pyexecutor/model_engine.py", - "PyTorchModelEngine", + "ModelEngine", { "prepare_cleanup", "shutdown_moe_load_balancer", @@ -677,7 +677,7 @@ def test_model_engine_prepare_failure_participates_in_unsafe_preflight(): events = [] engine_cls = _extract_class( _TORCH_ROOT / "pyexecutor/model_engine.py", - "PyTorchModelEngine", + "ModelEngine", { "prepare_cleanup", "shutdown_moe_load_balancer", @@ -717,7 +717,7 @@ def test_model_engine_finalizer_never_enters_collective_cleanup(): calls = [] engine_cls = _extract_class( _TORCH_ROOT / "pyexecutor/model_engine.py", - "PyTorchModelEngine", + "ModelEngine", {"__del__"}, {"logger": SimpleNamespace(warning=lambda *args: calls.append(("warning", args)))}, ) @@ -858,8 +858,6 @@ def async_worker_enabled(self): def async_worker_stop(self): self.stop_attempts += 1 events.append(("sampler_stop", self.stop_attempts)) - if self.stop_attempts == 1: - raise RuntimeError("sampler stop failed") class Engine: def release_executor_resources(self, *, release_model_owned_graphs): @@ -905,6 +903,8 @@ def __exit__(self, *args): executor.hang_detector = SimpleNamespace(detected=lambda: False) executor.worker_thread = SimpleNamespace(join=lambda: events.append("worker_join")) executor.kv_connector_manager = None + executor.enable_attention_dp = False + executor._kv_connector_failed = False executor.dist = SimpleNamespace(pp_size=1) executor._shutdown_sleep_wakeup_listeners = lambda: events.append("listener_shutdown") executor.encoder_launch_executor = None @@ -930,7 +930,7 @@ def terminal_cleanup(*, local_safe): worker.doing_shutdown = False worker.engine = executor - for message in ("pool delete failed", "sampler stop failed", "dwdp exit failed"): + for message in ("pool delete failed", "dwdp exit failed"): with pytest.raises(RuntimeError, match=message): worker.shutdown() assert executor.model_engine is engine @@ -947,14 +947,12 @@ def terminal_cleanup(*, local_safe): "profile_cleanup", "worker_join", "listener_shutdown", + ("sampler_stop", 1), "release_graphs", "resource_shutdown", ("pool_delete", 1), ("terminal_cleanup", False), ("pool_delete", 2), - ("sampler_stop", 1), - ("terminal_cleanup", False), - ("sampler_stop", 2), ("dwdp_exit", 1), ("terminal_cleanup", False), ("dwdp_exit", 2), From 361fbdc4d4ae42ac865ebb31508d724925f09ae1 Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 27 Oct 2026 19:50:53 -0700 Subject: [PATCH 51/59] [None][chore] Sync the in-tree CuTe DSL MegaMoE kernels with upstream 41a45039 Bring in the two upstream changes since the in-tree copy's base (5a2a6a59) that touch the files TensorRT-LLM carries: - 0c376e2f drops the fence.proxy.async.global after the Rubin FC2 STG epilogue store and removes the now unused PTX helper; - f29714ac reads the source token of an FC1 per-token activation scale with TokenSrcMetadata.unpack, matching the 16-bit token index of the routing metadata. TensorRT-LLM does not enable fc1_activation_per_token_scale. The other upstream merges since 5a2a6a59 change no file in this tree. The TensorRT-LLM adaptations are unchanged. Signed-off-by: jintaop --- .../communication/nvlink_domain/token_comm.py | 6 ++---- .../cutedsl_megamoe/helpers/ptx_helpers.py | 17 ----------------- .../mega/block_scaled_swap_ab_fc12_epilogue.py | 4 ---- 3 files changed, 2 insertions(+), 25 deletions(-) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm.py index ab052053ec9b..2fbb50436b53 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm.py @@ -772,10 +772,8 @@ def _router_push_metadata(self, smem_base: cute.Pointer) -> None: score = cute.arch.load(source_scores + route, cutlass.Float32) stg_f32(destination_scores_base + peer_offset + Int64(destination_position) * Int64(4), score) if cutlass.const_expr(self.fc1_activation_per_token_scale): - # The scale is per source token, not per route, so it is not sorted: gather it - # here by the token the metadata already names (low 32 bits, see - # TokenSrcMetadata) and push it into the same pool slot as the score. - source_token = Int32(Int64(metadata) & Int64(0xFFFFFFFF)) + # Gather the source token's scale into the same pool slot as its metadata. + source_token = TokenSrcMetadata.unpack(metadata).src_token activation_scale = self._router_activation_scales[source_token] stg_f32( destination_activation_scales_base + peer_offset + Int64(destination_position) * Int64(4), diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/ptx_helpers.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/ptx_helpers.py index 53192bb19bde..84b7cb973a34 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/ptx_helpers.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/ptx_helpers.py @@ -553,22 +553,6 @@ def cvt_rn_bf16x2_f32x2(acc2, *, loc: Optional[ir.Location] = None, ip: Optional return Int32(packed) -@dsl_user_op -def fence_proxy_async_global(*, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None) -> None: - """Order this thread's generic-proxy global accesses before its later async-proxy (TMA) accesses.""" - llvm.inline_asm( - None, - [], - "fence.proxy.async.global;", - "", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - @cute.jit def cvt_f32x4_to_f8x4_pack_i32(fp32x4: cute.Tensor, fp8_type, *, loc=None, ip=None) -> Int32: """Round four f32 lanes to the selected FP8 format and pack them into one i32.""" @@ -834,7 +818,6 @@ def stg_e8m0x8_from_f32( "stg_f32", "tma_load_1d", "fence_acq_rel_gpu", - "fence_proxy_async_global", "f32x2_pack", "fma_rn_f32x2_bf16x2", "cvt_rn_bf16x2_f32x2", diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py index fe912753b951..4d76ea7ffe3a 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py @@ -38,7 +38,6 @@ cp_async_bulk_s2g, cp_reduce_async_bulk_add_bf16_s2g, cvt_f32_to_fp8_to_f32, - fence_proxy_async_global, movmatrix_b16, red_add_relaxed_sys_v2_bf16x2, ) @@ -2614,9 +2613,6 @@ def fc2_stg_store_function( src_i = cute.make_tensor(copy_src[None, local_iter].iterator, single_copy_layout) dst_i = cute.make_tensor(dst_ptr, single_copy_layout) cute.copy(copy_atom_vec, cute.recast_tensor(src_i, cutlass.Int32), cute.recast_tensor(dst_i, cutlass.Int32)) - if cutlass.const_expr(epi.token_back_push_data): - # Generic-proxy pool rows are read back by the combine warps' TMA: fence before the fc2_done release. - fence_proxy_async_global() @cute.jit From 28d3cbd1db9107a691c71e6cfbc5bee59526a84e Mon Sep 17 00:00:00 2001 From: jintaop Date: Tue, 27 Oct 2026 20:28:07 -0700 Subject: [PATCH 52/59] [None][perf] Run HALO-Q as one general fused kernel Replace the FixedEp/FixedExperts/FixedHelpers specializations of the fused HALO-Q + in-switch copy kernel with one kernel that takes EP, experts and helpers at run time. The <8,384,4> and <4,384,4> instantiations, select_fused_kernel, the planner's <8,4>/<4,4> dispatch and the EP8 forms of exchange_counts and run_halo_q are removed. What the constants bought is kept with fixed-size structures that take runtime trip counts: - prepare_halo_m_register (EP <= 8, home <= 64, i.e. every EP8 geometry): one warp per owner sums its home masses over the source ranks into a two-round register array with clamped unconditional loads; only candidate masses reach shared memory. This was the EP8/4-helper path; every other geometry keeps the dense prepare. - plan_halo_m_register8: for EP <= 8 the hierarchy has at most two levels, so group selection is two full-mask sums and one ballot for every EP <= 8 and helper count (was EP8/4 helpers only; the per-level loop with per-group reductions is gone). - run_halo_q: the small-EP repair runs at most seven unrolled phases. - build_route_prefix: rank loops run in unrolled chunks of eight with clamped loads, and a chunk's loads precede its prefix stores. - color_and_publish_local: the free helper slot is found one mask word at a time with __ffs, so the kernel has no stack frame. - materialize_worker_hot_routes: each lane resolves four hot routes per batch, so their dependent load chains overlap. Plans, quotas, route prefixes, helper assignments, routed slots and status words are unchanged for every configuration. Signed-off-by: jintaop --- .../dynamicEplb/moeRebalanceHaloQKernel.cuh | 570 ++++++++---------- .../dynamicEplb/moeRebalanceTma.cu | 32 +- 2 files changed, 274 insertions(+), 328 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh index d71491dcfc13..3657a67671c6 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh @@ -41,6 +41,9 @@ constexpr int kMaxExperts = 384; constexpr int kMaxBroadcasts = kMaxExperts; constexpr int kHelperMaskWords = (kMaxBroadcasts + 31) / 32; constexpr int kMinGroup = 4; +/* Rank loops run in fully unrolled chunks of this many predicated ranks, so + * every load of a chunk is independent and in flight together. */ +constexpr int kRankChunk = 8; constexpr int kSymPayloadOffset = 64; using Epoch = std::uint32_t; @@ -358,11 +361,9 @@ __device__ bool stable_worker_ordinal_pass(int const* routes, int* ordinal, int* return true; } -template __device__ bool exchange_counts(unsigned long long const* peer_bases, int const* partial, int* shared, int ep, int expert_count, int local_rank, int ctas, Epoch epoch, unsigned long long spin_cycles, int* status) { - int const exchange_ep = FixedEp > 0 ? FixedEp : ep; int const tid = threadIdx.x; __shared__ int wait_failed; if (tid == 0) @@ -384,10 +385,10 @@ __device__ bool exchange_counts(unsigned long long const* peer_bases, int const* * complete row to the peer acquire. */ int const warp = tid >> 5; int const lane = tid & 31; - for (int destination = warp; destination < exchange_ep; destination += blockDim.x >> 5) + for (int destination = warp; destination < ep; destination += blockDim.x >> 5) { int* peer = reinterpret_cast(peer_bases[destination]); - int* row = peer + kSymPayloadOffset + (parity * exchange_ep + local_rank) * expert_count; + int* row = peer + kSymPayloadOffset + (parity * ep + local_rank) * expert_count; /* Keep scalar remote stores for the publication path; the local mirror * copy below uses vector transfers when alignment permits. */ for (int expert = lane; expert < expert_count; expert += 32) @@ -402,7 +403,7 @@ __device__ bool exchange_counts(unsigned long long const* peer_bases, int const* __syncthreads(); int* local = reinterpret_cast(peer_bases[local_rank]); - if (tid < exchange_ep) + if (tid < ep) { unsigned long long start = clock64(); Epoch observed = load_relaxed_system(local + tid); @@ -422,8 +423,8 @@ __device__ bool exchange_counts(unsigned long long const* peer_bases, int const* __syncthreads(); if (wait_failed) return false; - int const* payload = local + kSymPayloadOffset + parity * exchange_ep * expert_count; - int const total = exchange_ep * expert_count; + int const* payload = local + kSymPayloadOffset + parity * ep * expert_count; + int const total = ep * expert_count; /* Vectorize the local mirror copy when shape and alignment permit. */ if ((total & 3) == 0 && (reinterpret_cast(payload) & 15) == 0) { @@ -442,140 +443,145 @@ __device__ bool exchange_counts(unsigned long long const* peer_bases, int const* return true; } -__device__ void prepare_halo_m_sparse8(int const* counts, int* histogram, int* load, int* candidates, int experts) +/* Rank and home bounds of the register planner path. Every hierarchy level + * of EP <= 8 fits one 8-lane warp segment, and a home of at most 64 experts + * fits two 32-lane rounds. */ +constexpr int kRegisterEp = 8; +constexpr int kRegisterRounds = 2; + +__device__ __forceinline__ bool uses_register_prepare(int ep, int experts) { - int const tid = threadIdx.x; - int const warp = tid >> 5; - int const lane = tid & 31; - if (warp >= 8) - return; - int const home = experts / 8; - int local_index[2]; - int local_mass[2]; - int rank_load = 0; -#pragma unroll - for (int round = 0; round < 2; ++round) - { - int local = lane + 32 * round; - int mass = -1; - if (local < home) - { - mass = 0; - int expert = warp * home + local; -#pragma unroll - for (int source = 0; source < 8; ++source) - mass += counts[source * experts + expert]; - rank_load += mass; - } - local_index[round] = local; - local_mass[round] = mass; - } - rank_load = __reduce_add_sync(0xffffffffu, rank_load); - if (lane == 0) - load[warp] = rank_load; - for (int slot = 0; slot < 4; ++slot) + return ep <= kRegisterEp && experts / ep <= 32 * kRegisterRounds; +} + +/* Pick the owner's hot list from the lane-resident home masses: the heaviest + * remaining expert per slot, the lowest home index on a tie. Each lane holds + * home experts lane + 32 * round; entries outside the home carry -1. */ +template +__device__ __forceinline__ void select_owner_candidates(int (&local_mass)[Rounds], int* histogram, int* candidates, + int owner, int home, int candidate_count, bool write_mass) +{ + int const lane = threadIdx.x & 31; + for (int slot = 0; slot < candidate_count; ++slot) { int best_mass = -1; int best_local = INT_MAX; #pragma unroll - for (int round = 0; round < 2; ++round) + for (int round = 0; round < Rounds; ++round) { - int mass = local_mass[round]; - int local = local_index[round]; + int const mass = local_mass[round]; + int const local = lane + 32 * round; if (mass > best_mass || (mass == best_mass && local < best_local)) { best_mass = mass; best_local = local; } } - int warp_mass = __reduce_max_sync(0xffffffffu, best_mass); - int winner = __reduce_min_sync(0xffffffffu, best_mass == warp_mass ? best_local : INT_MAX); + int const warp_mass = __reduce_max_sync(0xffffffffu, best_mass); + int const winner = __reduce_min_sync(0xffffffffu, best_mass == warp_mass ? best_local : INT_MAX); if (lane == 0) { - int expert = warp * home + winner; - candidates[warp * 4 + slot] = expert; - histogram[expert] = warp_mass; + int const expert = owner * home + winner; + candidates[owner * candidate_count + slot] = expert; + if (write_mass) + histogram[expert] = warp_mass; } #pragma unroll - for (int round = 0; round < 2; ++round) - if (local_index[round] == winner) + for (int round = 0; round < Rounds; ++round) + if (lane + 32 * round == winner) local_mass[round] = -1; } - asm volatile("bar.sync 4, 256;" ::: "memory"); } -/* Build HALO-M's dense histogram, rank loads and per-owner hot lists with one - * warp per rank. The former thread-0 implementation dominated the kernel. */ -template +__device__ __forceinline__ void reset_owner_state( + int owner, int rank_load, int* load, int* recv_used, int* send_used, int* next_index, int* frozen) +{ + load[owner] = rank_load; + recv_used[owner] = 0; + send_used[owner] = 0; + next_index[owner] = 0; + frozen[owner] = 0; +} + +/* Rank loads and per-owner hot lists for uses_register_prepare(): one warp + * per owner sums its home masses over the source ranks straight into + * registers. Every load is unconditional (out-of-range sources and experts + * read a clamped valid address and contribute nothing), so all of them are in + * flight together. Only the candidates' masses reach histogram[]; the + * planners read no other entry. The owner warps, warp 0 among them, meet at + * a named barrier; the other warps leave at once. */ +__device__ void prepare_halo_m_register(int const* counts, int* histogram, int* load, int* candidates, int* recv_used, + int* send_used, int* next_index, int* frozen, int ep, int experts, int helpers) +{ + int const warp = threadIdx.x >> 5; + int const lane = threadIdx.x & 31; + if (warp >= ep) + return; + int const owner = warp; + int const home = experts / ep; + int local_mass[kRegisterRounds]; + int rank_load = 0; +#pragma unroll + for (int round = 0; round < kRegisterRounds; ++round) + { + int const local = lane + 32 * round; + bool const in_home = local < home; + int const expert = owner * home + imin(local, home - 1); + int mass = 0; +#pragma unroll + for (int source = 0; source < kRegisterEp; ++source) + { + int const count = counts[imin(source, ep - 1) * experts + expert]; + mass += source < ep ? count : 0; + } + rank_load += in_home ? mass : 0; + local_mass[round] = in_home ? mass : -1; + } + rank_load = __reduce_add_sync(0xffffffffu, rank_load); + if (lane == 0) + reset_owner_state(owner, rank_load, load, recv_used, send_used, next_index, frozen); + select_owner_candidates(local_mass, histogram, candidates, owner, home, imin(helpers, home), true); + asm volatile("bar.sync 4, %0;" ::"r"(ep * 32) : "memory"); +} + +/* Rounds of 32 home experts per lane; EP2 has the largest home, 192 experts. */ +constexpr int kMaxHomeRounds = (kMaxExperts / 2 + 31) / 32; + +/* Dense histogram, rank loads and per-owner hot lists for every other + * geometry, with one warp per owner. */ __device__ void prepare_halo_m(int const* counts, int* histogram, int* load, int* candidates, int* recv_used, int* send_used, int* next_index, int* frozen, int ep, int experts, int helpers) { - int const prepare_ep = FixedEp > 0 ? FixedEp : ep; - int const prepare_helpers = imin(FixedHelpers > 0 ? FixedHelpers : helpers, experts / prepare_ep); int const tid = threadIdx.x; int const warp = tid >> 5; int const lane = tid & 31; - int const home = experts / prepare_ep; + int const home = experts / ep; for (int expert = tid; expert < experts; expert += blockDim.x) { int mass = 0; #pragma unroll - for (int rank = 0; rank < prepare_ep; ++rank) + for (int rank = 0; rank < ep; ++rank) mass += counts[rank * experts + expert]; histogram[expert] = mass; } __syncthreads(); - for (int owner = warp; owner < prepare_ep; owner += blockDim.x >> 5) + for (int owner = warp; owner < ep; owner += blockDim.x >> 5) { int rank_load = 0; for (int local = lane; local < home; local += 32) rank_load += histogram[owner * home + local]; rank_load = __reduce_add_sync(0xffffffffu, rank_load); if (lane == 0) - { - load[owner] = rank_load; - recv_used[owner] = 0; - send_used[owner] = 0; - next_index[owner] = 0; - frozen[owner] = 0; - } - - /* EP2 is the largest home: 192 experts, six entries per lane. */ - constexpr int rounds = FixedEp == 8 ? 2 : 6; - int local_index[rounds]; - int local_mass[rounds]; + reset_owner_state(owner, rank_load, load, recv_used, send_used, next_index, frozen); + int local_mass[kMaxHomeRounds]; #pragma unroll - for (int round = 0; round < rounds; ++round) + for (int round = 0; round < kMaxHomeRounds; ++round) { - int local = lane + 32 * round; - local_index[round] = local; + int const local = lane + 32 * round; local_mass[round] = local < home ? histogram[owner * home + local] : -1; } - for (int slot = 0; slot < prepare_helpers; ++slot) - { - int best_mass = -1; - int best_local = INT_MAX; -#pragma unroll - for (int round = 0; round < rounds; ++round) - { - int mass = local_mass[round]; - int local = local_index[round]; - if (mass > best_mass || (mass == best_mass && local < best_local)) - { - best_mass = mass; - best_local = local; - } - } - int warp_mass = __reduce_max_sync(0xffffffffu, best_mass); - int winner = __reduce_min_sync(0xffffffffu, best_mass == warp_mass ? best_local : INT_MAX); - if (lane == 0) - candidates[owner * prepare_helpers + slot] = owner * home + winner; -#pragma unroll - for (int round = 0; round < rounds; ++round) - if (local_index[round] == winner) - local_mass[round] = -1; - } + select_owner_candidates(local_mass, histogram, candidates, owner, home, imin(helpers, home), false); } __syncthreads(); } @@ -583,12 +589,10 @@ __device__ void prepare_halo_m(int const* counts, int* histogram, int* load, int /* EP<=8 fast path: every participating lane owns one rank's complete * mutable HALO-M state. Group and destination decisions use only shuffles, * avoiding repeated shared-memory rank-state traffic and warp barriers. */ -template +template __device__ __forceinline__ int plan_halo_m_register8_impl( int const* histogram, int lane_load, int total, int* quota, int const* candidates, int* plan, int ep, int helpers) { - int const plan_ep = FixedEp > 0 ? FixedEp : ep; - int const plan_helpers = FixedHelpers > 0 ? FixedHelpers : helpers; int const lane = threadIdx.x & 31; constexpr unsigned mask = 0xffu; int lane_recv = 0; @@ -596,8 +600,8 @@ __device__ __forceinline__ int plan_halo_m_register8_impl( int lane_next = 0; int lane_frozen = 0; - int const target = total / plan_ep + (total % plan_ep != 0); - int const depth = hierarchy_depth(plan_ep); + int const target = total / ep + (total % ep != 0); + int const depth = hierarchy_depth(ep); int* p_expert = plan_field(plan, kPlanExpert); int* p_level = plan_field(plan, kPlanLevel); int* p_begin = plan_field(plan, kPlanGroupBegin); @@ -607,10 +611,10 @@ __device__ __forceinline__ int plan_halo_m_register8_impl( int broadcasts = 0; int trips = 0; - for (int guard = 0; guard < plan_ep * (plan_helpers + 2); ++guard) + for (int guard = 0; guard < ep * (helpers + 2); ++guard) { ++trips; - bool const eligible = lane < plan_ep && !lane_frozen && lane_load > target; + bool const eligible = lane < ep && !lane_frozen && lane_load > target; int owner, owner_load; if constexpr (NarrowKeys) { @@ -633,7 +637,7 @@ __device__ __forceinline__ int plan_halo_m_register8_impl( = __shfl_sync(mask, (static_cast(lane_send) << 4) | static_cast(lane_next), owner, 8); int const owner_send = static_cast(packed >> 4); int const owner_next = static_cast(packed & 15u); - if (owner_send >= plan_helpers || owner_next >= plan_helpers) + if (owner_send >= helpers || owner_next >= helpers) { if (lane == owner) lane_frozen = 1; @@ -642,11 +646,11 @@ __device__ __forceinline__ int plan_halo_m_register8_impl( /* owner and owner_next are warp-uniform here, so every lane can read * candidates[] and histogram[] directly. Both are written by - * prepare_halo_m_sparse8 before its bar.sync and never inside this + * the prepare step before its barrier and never inside this * loop, so every lane loads the same values and no broadcast from the * owner is needed. Only the owner advances lane_next, and it does so * before the mass test. */ - int const expert = candidates[owner * plan_helpers + owner_next]; + int const expert = candidates[owner * helpers + owner_next]; if (lane == owner) ++lane_next; int const mass = histogram[expert]; @@ -665,95 +669,39 @@ __device__ __forceinline__ int plan_halo_m_register8_impl( continue; } - int positive_target = imax(target, 1); - int need = imax(1, mass / positive_target + (mass % positive_target != 0)); - int wanted = depth - 1; - for (int level = depth - 1; level >= 0; --level) - { - if (plan_ep / (1 << level) >= need) - { - wanted = level; - break; - } - } - + /* EP <= 8 has at most two hierarchy levels: level 0 is every rank + * and, at EP8, level 1 is the two halves of four ranks. Two + * independent full-mask sums give the two half sums. Integer addition + * is exact, so zeroing the other half's addends yields each half sum + * exactly; the two reduces do not depend on each other, and neither + * redux.sync uses a per-lane membermask. The lower half wins a tie. */ + int const half = ep >> 1; + unsigned const rank_mask = (1u << ep) - 1u; + unsigned const low_mask = (1u << half) - 1u; + unsigned const high_mask = rank_mask & ~low_mask; + int const slack = lane < ep && lane != owner ? imax(0, target - lane_load) : 0; + int const room_low = __reduce_add_sync(mask, lane < half ? slack : 0); + int const room_high = __reduce_add_sync(mask, lane >= half ? slack : 0); + unsigned const receive_mask = __ballot_sync(mask, lane < ep && lane_recv < helpers); int chosen_level = -1; int chosen_begin = -1; int chosen_size = 0; - if constexpr (FixedEp == 8 && FixedHelpers == 4) + if (depth > 1 && mass <= half * target) { - int slack = lane != owner ? imax(0, target - lane_load) : 0; - /* Two independent full-mask sums give the two half sums. Integer - * addition is exact, so zeroing the other half's addends yields - * each half sum exactly; the two reduces do not depend on each - * other, and neither redux.sync uses a per-lane membermask. */ - int room0 = __reduce_add_sync(mask, lane < 4 ? slack : 0); - int room4 = __reduce_add_sync(mask, lane >= 4 ? slack : 0); - unsigned receive_mask = __ballot_sync(mask, lane_recv < plan_helpers); - bool valid0 = (receive_mask & 0x0fu) == 0x0fu && room0 > 0; - bool valid4 = (receive_mask & 0xf0u) == 0xf0u && room4 > 0; - bool want_half = mass <= 4 * target; - if (want_half && (valid0 || valid4)) + bool const valid_low = (receive_mask & low_mask) == low_mask && room_low > 0; + bool const valid_high = (receive_mask & high_mask) == high_mask && room_high > 0; + if (valid_low || valid_high) { chosen_level = 1; - chosen_begin = valid0 && (!valid4 || room0 >= room4) ? 0 : 4; - chosen_size = 4; - } - else if (receive_mask == mask && room0 + room4 > 0) - { - chosen_level = 0; - chosen_begin = 0; - chosen_size = 8; + chosen_begin = valid_low && (!valid_high || room_low >= room_high) ? 0 : half; + chosen_size = half; } } - else + if (chosen_level < 0 && (receive_mask & rank_mask) == rank_mask && room_low + room_high > 0) { - for (int level = wanted; level >= 0 && chosen_level < 0; --level) - { - int size = plan_ep / (1 << level); - bool active = lane < plan_ep; - int begin = (lane / size) * size; - unsigned group_mask = rank_group_mask(begin, size); - int local_room = active && lane != owner ? imax(0, target - lane_load) : 0; - int room = 0; - if (active) - room = __reduce_add_sync(group_mask, local_room); - unsigned receive_mask = __ballot_sync(mask, active && lane_recv < plan_helpers); - bool group_leader = active && lane == begin && (receive_mask & group_mask) == group_mask; - unsigned leader_mask = __ballot_sync(mask, group_leader); - int best_room = -1; - int best_begin = INT_MAX; - if (leader_mask) - { - int const first_leader = __ffs(leader_mask) - 1; - if constexpr (NarrowKeys) - { - unsigned key = 0; - if (group_leader) - key = __reduce_max_sync( - leader_mask, (static_cast(room) << 4) | (15u - static_cast(begin))); - key = __shfl_sync(mask, key, first_leader, 8); - best_room = static_cast(key >> 4); - best_begin = 15 - static_cast(key & 15u); - } - else - { - if (group_leader) - { - best_room = __reduce_max_sync(leader_mask, room); - best_begin = __reduce_min_sync(leader_mask, room == best_room ? begin : INT_MAX); - } - best_room = __shfl_sync(mask, best_room, first_leader, 8); - best_begin = __shfl_sync(mask, best_begin, first_leader, 8); - } - } - if (best_begin != INT_MAX && best_room > 0) - { - chosen_level = level; - chosen_begin = best_begin; - chosen_size = size; - } - } + chosen_level = 0; + chosen_begin = 0; + chosen_size = ep; } if (chosen_level < 0) { @@ -819,8 +767,8 @@ __device__ __forceinline__ int plan_halo_m_register8_impl( p_owner[b] = owner; p_helper[b] = -1; } - if (lane < plan_ep) - quota[b * plan_ep + lane] = lane_quota; + if (lane < ep) + quota[b * ep + lane] = lane_quota; bool in_group = lane >= chosen_begin && lane < chosen_begin + chosen_size; if (in_group) @@ -839,22 +787,18 @@ __device__ __forceinline__ int plan_halo_m_register8_impl( } /* Select the key representation once, from the actual complete rank load. - * This guard keeps owner/destination loads and group-room keys representable. + * This guard keeps the owner and destination load keys representable. * Keep the wide path for large valid int32 inputs without putting a runtime * representation branch on each iteration of either planning loop. */ -template __device__ __forceinline__ int plan_halo_m_register8( int const* histogram, int const* load, int* quota, int const* candidates, int* plan, int ep, int helpers) { - int const plan_ep = FixedEp > 0 ? FixedEp : ep; int const lane = threadIdx.x & 31; - int const lane_load = lane < plan_ep ? load[lane] : 0; + int const lane_load = lane < ep ? load[lane] : 0; int const total = __reduce_add_sync(0xffu, lane_load); if (total <= 0x0fffffff) - return plan_halo_m_register8_impl( - histogram, lane_load, total, quota, candidates, plan, ep, helpers); - return plan_halo_m_register8_impl( - histogram, lane_load, total, quota, candidates, plan, ep, helpers); + return plan_halo_m_register8_impl(histogram, lane_load, total, quota, candidates, plan, ep, helpers); + return plan_halo_m_register8_impl(histogram, lane_load, total, quota, candidates, plan, ep, helpers); } /* Exact HALO-M state machine executed by warp 0. Each lane owns one rank; @@ -1224,8 +1168,8 @@ __device__ __forceinline__ void halo_q_direction8( * `b`, `hosts` and `expert` are the lane's own broadcast identity and do not * change between phases or rounds, so the caller hoists them. Reading them * here would put a four-deep dependent load chain (order[lane] -> b -> host_mask[b] - * and the GLOBAL p_expert[b] -> expert -> counts[...]) at the head of all - * fourteen phases. */ + * and the GLOBAL p_expert[b] -> expert -> counts[...]) at the head of every + * phase. */ __device__ __forceinline__ void halo_q_pair8( int left, int right, int b, int hosts, int expert, int const* counts, int* quota, int ep, int experts) { @@ -1312,11 +1256,9 @@ __device__ __forceinline__ void halo_q_build_order_and_masks(int const* p_expert host_mask[lane] = static_cast(rank_group_mask(p_begin[lane], p_size[lane]) | (1u << p_owner[lane])); } -template __device__ void run_halo_q(int const* counts, int* quota, int* plan, int broadcasts, int ep, int experts, int* ring, int* order, int* host_mask) { - int const plan_ep = FixedEp > 0 ? FixedEp : ep; int const tid = threadIdx.x; int const warp = tid >> 5; int const lane = tid & 31; @@ -1324,31 +1266,37 @@ __device__ void run_halo_q(int const* counts, int* quota, int* plan, int broadca int* p_begin = plan_field(plan, kPlanGroupBegin); int* p_size = plan_field(plan, kPlanGroupSize); int* p_owner = plan_field(plan, kPlanOwner); - if (plan_ep <= 8 && broadcasts <= 8) + if (ep <= kRegisterEp && broadcasts <= 8) { - if (warp < plan_ep / 2) + if (warp < ep / 2) { if (warp == 0) halo_q_build_order_and_masks(p_expert, p_begin, p_size, p_owner, order, host_mask, broadcasts); - halo_q_active_barrier((plan_ep / 2) * 32); + halo_q_active_barrier((ep / 2) * 32); /* Loop-invariant broadcast identity, read once after the barrier - * that publishes order[] and host_mask[]. All four participating + * that publishes order[] and host_mask[]. All ep / 2 participating * warps read the same lane-indexed entries. */ int const my_b = lane < broadcasts ? order[lane] : -1; int const my_hosts = my_b >= 0 ? host_mask[my_b] : 0; int const my_expert = my_b >= 0 ? p_expert[my_b] : -1; if (broadcasts >= 2) { - int const ring_size = plan_ep - 1; + int const ring_size = ep - 1; /* Run two repair rounds unconditionally. A phase with no required * repair is a no-op, so the plan equals running only the required * repairs. */ int const repair_rounds = 2; for (int round = 0; round < repair_rounds; ++round) { - for (int phase = 0; phase < plan_ep - 1; ++phase) + /* At most kRegisterEp - 1 phases, unrolled; the guard is + * uniform across the block, so every pair warp meets + * every barrier it executes. */ +#pragma unroll + for (int phase = 0; phase < kRegisterEp - 1; ++phase) { - int right_index = warp == 0 ? 0 : plan_ep - 1 - warp; + if (phase >= ep - 1) + break; + int right_index = warp == 0 ? 0 : ep - 1 - warp; int right_position = right_index - phase; if (right_position < 0) right_position += ring_size; @@ -1361,22 +1309,16 @@ __device__ void run_halo_q(int const* counts, int* quota, int* plan, int broadca left_position += ring_size; left = 1 + left_position; } - halo_q_pair8(left, right, my_b, my_hosts, my_expert, counts, quota, plan_ep, experts); - if constexpr (FixedEp == 8) - { - /* Must stay unconditional. Omitting it after a - * round's last phase races twice: warp 0 enters the - * next round's phase 0 while warp 3 still writes the - * same quota[], and after the final phase warps 1-3 - * leave run_halo_q while warp 0 -- which only meets - * warp 8 at scheduler_tail_barrier -- walks into - * build_route_prefix and reads every quota[]. */ - halo_q_active_barrier(128); - } - else - { - halo_q_active_barrier((plan_ep / 2) * 32); - } + halo_q_pair8(left, right, my_b, my_hosts, my_expert, counts, quota, ep, experts); + /* Must stay unconditional. Omitting it after a + * round's last phase races twice: warp 0 enters the + * next round's phase 0 while the last pair warp still + * writes the same quota[], and after the final phase + * the other pair warps leave run_halo_q while warp 0 + * -- which only meets warp 8 at scheduler_tail_barrier + * -- walks into build_route_prefix and reads every + * quota[]. */ + halo_q_active_barrier((ep / 2) * 32); } } } @@ -1385,7 +1327,7 @@ __device__ void run_halo_q(int const* counts, int* quota, int* plan, int broadca } if (warp == 0) halo_q_build_order_and_masks(p_expert, p_begin, p_size, p_owner, order, host_mask, broadcasts); - if (tid < plan_ep - 1) + if (tid < ep - 1) ring[tid] = tid + 1; __syncthreads(); if (broadcasts < 2) @@ -1393,9 +1335,9 @@ __device__ void run_halo_q(int const* counts, int* quota, int* plan, int broadca for (int round = 0; round < 2; ++round) { - for (int phase = 0; phase < plan_ep - 1; ++phase) + for (int phase = 0; phase < ep - 1; ++phase) { - if (warp < plan_ep / 2) + if (warp < ep / 2) { int left, right; if (warp == 0) @@ -1406,16 +1348,16 @@ __device__ void run_halo_q(int const* counts, int* quota, int* plan, int broadca else { left = ring[warp]; - right = ring[plan_ep - 1 - warp]; + right = ring[ep - 1 - warp]; } - halo_q_repair(left, right, broadcasts, order, host_mask, counts, p_expert, quota, plan_ep, experts); - halo_q_repair(right, left, broadcasts, order, host_mask, counts, p_expert, quota, plan_ep, experts); + halo_q_repair(left, right, broadcasts, order, host_mask, counts, p_expert, quota, ep, experts); + halo_q_repair(right, left, broadcasts, order, host_mask, counts, p_expert, quota, ep, experts); } __syncthreads(); if (tid == 0) { - int last = ring[plan_ep - 2]; - for (int index = plan_ep - 2; index > 0; --index) + int last = ring[ep - 2]; + for (int index = ep - 2; index > 0; --index) ring[index] = ring[index - 1]; ring[0] = last; } @@ -1445,17 +1387,20 @@ __device__ void color_and_publish_local(int* plan, int broadcasts, int ep, int h { if (p_level[b] != level) continue; - unsigned unavailable[kHelperMaskWords] = {0}; - for (int rank = p_begin[b]; rank < p_begin[b] + p_size[b]; ++rank) - for (int word = 0; word < words; ++word) - unavailable[word] |= static_cast(used[rank * kHelperMaskWords + word]); + /* The lowest slot free on every group rank, one mask word at a + * time, so the scan keeps no per-word array. */ int helper = -1; - for (int slot = 0; slot < active_helpers; ++slot) - if (!(unavailable[slot / 32] & (1u << (slot % 32)))) - { - helper = slot; - break; - } + for (int word = 0; word < words && helper < 0; ++word) + { + unsigned unavailable = 0; + for (int rank = p_begin[b]; rank < p_begin[b] + p_size[b]; ++rank) + unavailable |= static_cast(used[rank * kHelperMaskWords + word]); + int const slots = imin(32, active_helpers - word * 32); + unsigned const valid = slots == 32 ? 0xffffffffu : (1u << slots) - 1u; + unsigned const open_slots = ~unavailable & valid; + if (open_slots) + helper = word * 32 + __ffs(open_slots) - 1; + } if (helper < 0) { record_status(status + 4, epoch); @@ -1493,10 +1438,18 @@ __device__ void build_route_prefix(int const* counts, int const* quota, int* pre { int expert = p_expert[b]; int source_before = 0; - for (int rank = 0; rank < local_rank; ++rank) + /* Ranks past the bound load a clamped valid entry and add nothing, + * so every load of a chunk is unconditional. */ + for (int base = 0; base < local_rank; base += kRankChunk) { - int diagonal = imin(counts[rank * experts + expert], quota[b * ep + rank]); - source_before += counts[rank * experts + expert] - diagonal; +#pragma unroll + for (int step = 0; step < kRankChunk; ++step) + { + int const rank = imin(base + step, local_rank - 1); + int const count = counts[rank * experts + expert]; + int const diagonal = imin(count, quota[b * ep + rank]); + source_before += base + step < local_rank ? count - diagonal : 0; + } } int source_count = counts[local_rank * experts + expert]; int source_diagonal = imin(source_count, quota[b * ep + local_rank]); @@ -1508,16 +1461,28 @@ __device__ void build_route_prefix(int const* counts, int const* quota, int* pre * although route_destination never reads it: it costs one store and * leaves no stale entry below ep. */ int destination_before = 0; - for (int destination = 0; destination < ep; ++destination) + for (int base = 0; base < ep; base += kRankChunk) { - int overlap = destination_before - source_before; - overlap = imax(0, imin(residual_supply, overlap)); - int diagonal_prefix = local_rank < destination ? source_diagonal : 0; - prefix[b * (ep + 1) + destination] = diagonal_prefix + overlap; - if (destination + 1 < ep) + /* Load the whole chunk before the first prefix store, which the + * compiler cannot order after loads it might alias. */ + int destination_supply[kRankChunk]; +#pragma unroll + for (int step = 0; step < kRankChunk; ++step) { - int diagonal = imin(counts[destination * experts + expert], quota[b * ep + destination]); - destination_before += quota[b * ep + destination] - diagonal; + int const destination = imin(base + step, ep - 1); + int const destination_quota = quota[b * ep + destination]; + destination_supply[step] + = destination_quota - imin(counts[destination * experts + expert], destination_quota); + } +#pragma unroll + for (int step = 0; step < kRankChunk; ++step) + { + int const destination = base + step; + int const overlap = imax(0, imin(residual_supply, destination_before - source_before)); + int const diagonal_prefix = local_rank < destination ? source_diagonal : 0; + if (destination < ep) + prefix[b * (ep + 1) + destination] = diagonal_prefix + overlap; + destination_before += destination + 1 < ep ? destination_supply[step] : 0; } } } @@ -1676,28 +1641,51 @@ __device__ void materialize_worker_hot_routes(int const* routes, int* output, in int const end = imin(worker_end, begin + warp_span); int hot_count = shared[warp * bins + experts]; hot_count = imin(hot_count, imax(0, end - begin)); - for (int offset = lane; offset < hot_count; offset += 32) + /* Each lane resolves kHotBatch routes at once: every dependent load of a + * stage is issued for the whole batch before the next stage consumes it. + * Batch members are distinct routes, so the order of their stores does + * not matter. */ + constexpr int kHotBatch = 4; + for (int first = lane; first < hot_count; first += 32 * kHotBatch) { - int index = hot_indices[begin + offset]; - int expert = routes[index]; - int b = expert_lut[expert]; - int ordinal = output[index] + shared[warp * bins + expert]; - int const* prefix = route_prefix + b * (ep + 1); - int destination = route_destination(prefix, ordinal, ep); - int slot = -1; - if (group_contains(destination, p_begin[b], p_size[b])) + int index[kHotBatch]; + int expert[kHotBatch]; + int ordinal[kHotBatch]; +#pragma unroll + for (int step = 0; step < kHotBatch; ++step) { - slot = destination * local_slots + home + p_helper[b]; + int const offset = first + 32 * step; + index[step] = offset < hot_count ? hot_indices[begin + offset] : -1; } - else if (destination == p_owner[b]) +#pragma unroll + for (int step = 0; step < kHotBatch; ++step) { - slot = destination * local_slots + expert - p_owner[b] * home; + expert[step] = index[step] >= 0 ? routes[index[step]] : 0; + ordinal[step] = index[step] >= 0 ? output[index[step]] : 0; } - else +#pragma unroll + for (int step = 0; step < kHotBatch; ++step) { - record_status(status + 5, epoch); + if (index[step] < 0) + continue; + int b = expert_lut[expert[step]]; + int const* prefix = route_prefix + b * (ep + 1); + int destination = route_destination(prefix, ordinal[step] + shared[warp * bins + expert[step]], ep); + int slot = -1; + if (group_contains(destination, p_begin[b], p_size[b])) + { + slot = destination * local_slots + home + p_helper[b]; + } + else if (destination == p_owner[b]) + { + slot = destination * local_slots + expert[step] - p_owner[b] * home; + } + else + { + record_status(status + 5, epoch); + } + output[index[step]] = slot; } - output[index] = slot; } } @@ -1766,15 +1754,11 @@ __device__ void materialize_worker_routes(int const* routes, int* output, int co /* One scheduler generation on a grid of ctas resident CTAs of 512 threads: * CTA 0 plans and publishes the helper placement, the other CTAs write routes. * The fused scheduler + in-switch copy kernel runs it as its first phase. */ -template __device__ __forceinline__ void halo_q_scheduler_body(int const* routes, int* out_slots, int* out_ids, int* out_levels, int* out_owners, unsigned long long const* peer_bases, int* status, int* partial, int* route_aux, int* grid_sync, - int* plan, int* route_prefix, int runtime_ep, int runtime_experts, int runtime_helpers, int local_rank, - int route_capacity, int ctas, unsigned long long spin_cycles, int capacity, int route_count) + int* plan, int* route_prefix, int ep, int experts, int helpers, int local_rank, int route_capacity, int ctas, + unsigned long long spin_cycles, int capacity, int route_count) { - int const ep = FixedEp > 0 ? FixedEp : runtime_ep; - int const experts = FixedExperts > 0 ? FixedExperts : runtime_experts; - int const helpers = FixedHelpers > 0 ? FixedHelpers : runtime_helpers; // Only the caller's valid prefix is readable. Keep the public output at // its configured capacity, clearing stale routes inside this same kernel. // Empty ranks still execute every collective; no route load occurs at N=0. @@ -1814,26 +1798,14 @@ __device__ __forceinline__ void halo_q_scheduler_body(int const* routes, int* ou offset += capacity; int* host_mask = shared + offset; - bool exchange_ok; - if (ep == 8) - exchange_ok = exchange_counts<8>( - peer_bases, partial, shared, ep, experts, local_rank, ctas, epoch, spin_cycles, status); - else - exchange_ok = exchange_counts<0>( - peer_bases, partial, shared, ep, experts, local_rank, ctas, epoch, spin_cycles, status); - if (!exchange_ok) + if (!exchange_counts(peer_bases, partial, shared, ep, experts, local_rank, ctas, epoch, spin_cycles, status)) return; - /* EP8 keeps only its candidate masses; other EPs retain the dense - * histogram required by the generic planner path. */ - if (ep == 8 && helpers == 4 && experts >= 32) - { - prepare_halo_m_sparse8(counts, histogram, load, selected_flags, experts); - } + if (uses_register_prepare(ep, experts)) + prepare_halo_m_register(counts, histogram, load, selected_flags, recv_used, send_used, next_index, frozen, + ep, experts, helpers); else - { - prepare_halo_m<0, 0>(counts, histogram, load, selected_flags, recv_used, send_used, next_index, frozen, ep, + prepare_halo_m(counts, histogram, load, selected_flags, recv_used, send_used, next_index, frozen, ep, experts, helpers); - } int broadcasts = 0; if (threadIdx.x < (ep <= 8 ? 8 : 32)) @@ -1841,18 +1813,7 @@ __device__ __forceinline__ void halo_q_scheduler_body(int const* routes, int* ou // These short-circuit bounds keep helpers*ep <= 120. if (ep <= 8 && helpers <= 15 && helpers * ep <= experts) { - if (ep == 8 && helpers == 4) - { - broadcasts = plan_halo_m_register8<8, 4>(histogram, load, quota, selected_flags, plan, ep, helpers); - } - else if (ep == 4 && helpers == 4) - { - broadcasts = plan_halo_m_register8<4, 4>(histogram, load, quota, selected_flags, plan, ep, helpers); - } - else - { - broadcasts = plan_halo_m_register8<0, 0>(histogram, load, quota, selected_flags, plan, ep, helpers); - } + broadcasts = plan_halo_m_register8(histogram, load, quota, selected_flags, plan, ep, helpers); } else { @@ -1876,10 +1837,7 @@ __device__ __forceinline__ void halo_q_scheduler_body(int const* routes, int* ou if (threadIdx.x == 8 * 32) color_and_publish_local( plan, broadcasts, ep, helpers, local_rank, recv_used, out_ids, out_levels, out_owners, status, epoch); - if (ep == 8) - run_halo_q<8>(counts, quota, plan, broadcasts, ep, experts, ring, order, host_mask); - else - run_halo_q<0>(counts, quota, plan, broadcasts, ep, experts, ring, order, host_mask); + run_halo_q(counts, quota, plan, broadcasts, ep, experts, ring, order, host_mask); if (ctas > 1) { int const warp = threadIdx.x >> 5; diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu index adec20d46c8a..cbd487168663 100644 --- a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu @@ -744,17 +744,15 @@ struct FusedSchedulerArgs * consumers run in the default one, so the release is system-scope. Each CTA * also waits for route_target before it reads the plan, which CTA 0 finishes * last, with the scheduler's spin bound. */ -template __global__ void __launch_bounds__(512, 1) halo_q_tma_fused_kernel(FusedSchedulerArgs const scheduler, GpuDirectStorage const direct, int warps, int total_slots, unsigned* completed_ctas, uint64_t flag_mc, uint64_t generation, unsigned long long* route_epoch, unsigned long long route_target) { - tensorrt_llm::kernels::halo_q_scheduler_body(scheduler.routes, - scheduler.out_slots, scheduler.out_ids, scheduler.out_levels, scheduler.out_owners, scheduler.peer_bases, - scheduler.status, scheduler.partial, scheduler.route_aux, scheduler.grid_sync, scheduler.plan, - scheduler.route_prefix, scheduler.ep, scheduler.experts, scheduler.helpers, scheduler.local_rank, - scheduler.route_capacity, static_cast(gridDim.x), scheduler.spin_cycles, scheduler.capacity, - scheduler.route_count); + tensorrt_llm::kernels::halo_q_scheduler_body(scheduler.routes, scheduler.out_slots, scheduler.out_ids, + scheduler.out_levels, scheduler.out_owners, scheduler.peer_bases, scheduler.status, scheduler.partial, + scheduler.route_aux, scheduler.grid_sync, scheduler.plan, scheduler.route_prefix, scheduler.ep, + scheduler.experts, scheduler.helpers, scheduler.local_rank, scheduler.route_capacity, + static_cast(gridDim.x), scheduler.spin_cycles, scheduler.capacity, scheduler.route_count); // Orders every thread's scheduler writes before thread 0's release. __syncthreads(); if (threadIdx.x == 0) @@ -778,16 +776,6 @@ __global__ void __launch_bounds__(512, 1) halo_q_tma_fused_kernel(FusedScheduler tma_copy_gpu_direct_body(direct, warps, total_slots, completed_ctas, flag_mc, generation); } -// HALO-Q is specialized for 384 experts with 4 helpers at EP4 and EP8. -void const* select_fused_kernel(MegamoeTmaGpuPlanConfig const& config) -{ - if (config.global_experts == 384 && config.helper_count == 4 && config.world == 8) - return reinterpret_cast(&halo_q_tma_fused_kernel<8, 384, 4>); - if (config.global_experts == 384 && config.helper_count == 4 && config.world == 4) - return reinterpret_cast(&halo_q_tma_fused_kernel<4, 384, 4>); - return reinterpret_cast(&halo_q_tma_fused_kernel<0, 0, 0>); -} - } // namespace struct MegamoeTmaCopyState @@ -801,7 +789,7 @@ struct MegamoeTmaCopyState uint64_t* direct_tables = nullptr; bool direct_configured = false; bool direct_bound = false; - // The fused kernel of the bound geometry, prepared by megamoe_tma_copy_bind_gpu_direct. + // The fused kernel, prepared by megamoe_tma_copy_bind_gpu_direct. void const* fused_kernel = nullptr; int fused_total_slots = 0; int fused_dynamic_bytes = 0; @@ -825,7 +813,7 @@ cudaError_t prepare_fused_launch(MegamoeTmaCopyState& state) { auto const& c = state.config; auto const& plan = state.direct.config; - void const* kernel = select_fused_kernel(plan); + void const* kernel = reinterpret_cast(&halo_q_tma_fused_kernel); cudaFuncAttributes attributes{}; cudaError_t error = cudaFuncGetAttributes(&attributes, kernel); if (error != cudaSuccess) @@ -985,9 +973,9 @@ extern "C" int megamoe_tma_copy_create(uint64_t max_segments, int sms, int warps &config.device_shared_bytes_per_sm, cudaDevAttrMaxSharedMemoryPerMultiprocessor, config.device); if (error != cudaSuccess) return fail(error); - // The generic fused kernel bounds the geometry; binding prepares the - // specialization of the bound plan geometry. - void const* kernel = reinterpret_cast(&halo_q_tma_fused_kernel<0, 0, 0>); + // The fused kernel's static shared memory bounds the copy slots; binding + // prepares its launch for the bound plan geometry. + void const* kernel = reinterpret_cast(&halo_q_tma_fused_kernel); cudaFuncAttributes attributes{}; error = cudaFuncGetAttributes(&attributes, kernel); if (error != cudaSuccess) From 47bed43b268d1ae4644eb8e8ca531116b59175d8 Mon Sep 17 00:00:00 2001 From: jintaop Date: Thu, 29 Oct 2026 02:23:39 -0700 Subject: [PATCH 53/59] [None][chore] Drop the trtllm-bench GC switch and the duplicate ruff exclude Keep trtllm-bench and pyproject.toml as on main: - remove TRTLLM_BENCH_DISABLE_GC and its measurement-run GC context from trtllm-bench throughput and latency, and its paragraph in perf-benchmarking.md; benchmark harnesses that need the benchmark process's GC off can do so outside trtllm-bench; - remove the ruff exclude for the vendored MegaMoE kernel and scheduler trees from pyproject.toml: the top-level exclude in .pre-commit-config.yaml already keeps every pre-commit hook, ruff included, off those trees. Signed-off-by: jintaop --- .../developer-guide/perf-benchmarking.md | 11 +++---- pyproject.toml | 3 -- tensorrt_llm/bench/benchmark/low_latency.py | 5 ++- tensorrt_llm/bench/benchmark/throughput.py | 5 ++- tensorrt_llm/bench/benchmark/utils/general.py | 31 +------------------ 5 files changed, 9 insertions(+), 46 deletions(-) diff --git a/docs/source/developer-guide/perf-benchmarking.md b/docs/source/developer-guide/perf-benchmarking.md index 9738e4c0ceb0..66f280350890 100644 --- a/docs/source/developer-guide/perf-benchmarking.md +++ b/docs/source/developer-guide/perf-benchmarking.md @@ -256,13 +256,10 @@ trtllm-bench --model meta-llama/Llama-3.1-8B \ With many short requests, for example context-only runs at high concurrency, rank 0 can be bound by host work, and a pause there makes every other rank wait -at the next collective. With `TRTLLM_BENCH_DISABLE_GC=1`, the `throughput` and -`latency` commands collect once after warmup, freeze the surviving objects, and -disable Python's cyclic garbage collector for the measured run, so a full -collection over the in-flight requests cannot pause the benchmark process. -`TRTLLM_WORKER_DISABLE_GC=1` does the same in the worker processes. Memory held -in reference cycles is then not freed until the run ends, so set the variables -identically in every run you compare. +at the next collective. `TRTLLM_WORKER_DISABLE_GC=1` disables Python's cyclic +garbage collector in the worker processes, so a full collection over the +in-flight requests cannot pause a rank. Memory held in reference cycles is then +not freed until the run ends, so set it identically in every run you compare. With `TRTLLM_PROXY_BATCH_REQUESTS=1`, the executor proxy pickles each request on the submitting thread and a background thread sends each burst to the rank-0 diff --git a/pyproject.toml b/pyproject.toml index d8264fc74a57..bbbe3144ae5f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,9 +49,6 @@ exclude = [ "tensorrt_llm/_torch/moe/kernels/rubin/manual_mma_128dp.py", "tensorrt_llm/_torch/moe/kernels/rubin/rubin_contiguous_grouped_blockscaled_gemm_fused_fc12.py", "tensorrt_llm/grpc/openengine/_generated/**", - # Keep vendored MegaMoE sources in upstream formatting. - "tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/**", - "tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/**", ] diff --git a/tensorrt_llm/bench/benchmark/low_latency.py b/tensorrt_llm/bench/benchmark/low_latency.py index 96a142ad2932..25fa2ee9e828 100644 --- a/tensorrt_llm/bench/benchmark/low_latency.py +++ b/tensorrt_llm/bench/benchmark/low_latency.py @@ -27,8 +27,7 @@ generate_json_report, get_general_cli_options, get_llm) from tensorrt_llm.bench.benchmark.utils.asynchronous import async_benchmark -from tensorrt_llm.bench.benchmark.utils.general import (generate_warmup_dataset, - measurement_gc) +from tensorrt_llm.bench.benchmark.utils.general import generate_warmup_dataset from tensorrt_llm.bench.dataclasses.configuration import RuntimeConfig from tensorrt_llm.bench.dataclasses.general import BenchmarkEnvironment from tensorrt_llm.bench.dataclasses.reporting import ReportUtility @@ -340,7 +339,7 @@ def latency_command( logger.info("Warmup done.") iteration_writer = options.iteration_writer - with iteration_writer.capture(), measurement_gc(): + with iteration_writer.capture(): statistics = asyncio.run( async_benchmark(llm, sampling_params, diff --git a/tensorrt_llm/bench/benchmark/throughput.py b/tensorrt_llm/bench/benchmark/throughput.py index 22a8f55fb53e..a28d6e00a2b2 100755 --- a/tensorrt_llm/bench/benchmark/throughput.py +++ b/tensorrt_llm/bench/benchmark/throughput.py @@ -36,8 +36,7 @@ get_settings) # isort: on from tensorrt_llm.bench.benchmark.utils.general import ( - generate_warmup_dataset, measurement_gc, - update_sampler_args_with_extra_options) + generate_warmup_dataset, update_sampler_args_with_extra_options) from tensorrt_llm.bench.dataclasses.configuration import RuntimeConfig from tensorrt_llm.bench.dataclasses.general import BenchmarkEnvironment from tensorrt_llm.bench.dataclasses.reporting import ReportUtility @@ -495,7 +494,7 @@ def throughput_command( logger.info("Warmup done.") iteration_writer = options.iteration_writer - with iteration_writer.capture(), measurement_gc(): + with iteration_writer.capture(): statistics = asyncio.run( async_benchmark(llm, sampling_params, diff --git a/tensorrt_llm/bench/benchmark/utils/general.py b/tensorrt_llm/bench/benchmark/utils/general.py index e98f90337c48..cf567b89a68b 100755 --- a/tensorrt_llm/bench/benchmark/utils/general.py +++ b/tensorrt_llm/bench/benchmark/utils/general.py @@ -1,11 +1,8 @@ from __future__ import annotations -import gc -import os -from contextlib import contextmanager from pathlib import Path from random import choices, shuffle -from typing import Dict, Iterator, List, Union +from typing import Dict, List, Union import yaml @@ -171,32 +168,6 @@ def get_settings(params: dict, dataset_metadata: DatasetMetadata, model: str, } -@contextmanager -def measurement_gc() -> Iterator[None]: - """Keep Python's cyclic GC out of the measured run if TRTLLM_BENCH_DISABLE_GC=1. - - The benchmark process submits every request and handles every response. A - full collection over its in-flight and finished requests can stop it long - enough for the executors to run out of requests. This collects once, - freezes the surviving objects and disables the GC until the measured run - ends, as TRTLLM_SERVER_DISABLE_GC and TRTLLM_WORKER_DISABLE_GC disable it - in the server and the workers. - """ - if os.getenv("TRTLLM_BENCH_DISABLE_GC", "0") != "1": - yield - return - gc.collect() - gc.freeze() - gc.disable() - logger.info("Python GC disabled for the measured run " - "(TRTLLM_BENCH_DISABLE_GC=1).") - try: - yield - finally: - gc.enable() - gc.unfreeze() - - def generate_warmup_dataset(requests, steps) -> List[InferenceRequest]: """Warm up the benchmarker.""" warm_up_dataset = choices(requests, k=steps) From fcc5bd22858bc36fbbb72607e68bf7e351c7196a Mon Sep 17 00:00:00 2001 From: jintaop Date: Thu, 29 Oct 2026 05:07:56 -0700 Subject: [PATCH 54/59] [None][test] Give the chunked-forward MoE test mock the perfect-router hook The external-comm scheduler now passes the routed expert ids through MoE._maybe_get_perfect_router_expert_ids. The SimpleNamespace layer that the B12x and fresh-output chunked-forward tests (main #21057) drive the scheduler with did not define it; give it the identity hook, which is what MoE does without ENABLE_PERFECT_ROUTER=expert_ids. Signed-off-by: jintaop --- tests/unittest/_torch/moe/test_moe_backend.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index b239ebad4aa6..e9700906bbe6 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -6397,6 +6397,7 @@ def route(router_logits, input_ids=None): _load_balancer_start_wait_gpu_stage=no_op, _load_balancer_start_set_cpu_stage=no_op, _load_balancer_done_set_cpu_stage=no_op, + _maybe_get_perfect_router_expert_ids=lambda token_selected_experts: token_selected_experts, ) From 4b919ae8c48b3db0ebe9ac67003a632bc8211a61 Mon Sep 17 00:00:00 2001 From: jintaop Date: Thu, 29 Oct 2026 07:39:12 -0700 Subject: [PATCH 55/59] [None][fix] Default DeepseekV4MoE's per-layer dynamic EPLB flag on the class DeepseekV4MoE.forward() reads _per_layer_dynamic_eplb_enabled, which __init__ sets. main's test_empty_decoder_moe_runs_expert_communication_ without_local_gemms builds the module with __new__ and failed with an AttributeError. A class-level False covers such instances; every module built through __init__ still sets its own value, so behavior is unchanged. Signed-off-by: jintaop --- tensorrt_llm/_torch/models/modeling_deepseekv4.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 4d6c29ba6b9c..bb8308465e6f 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -1625,6 +1625,9 @@ def get_experts_per_token(self): class DeepseekV4MoE(nn.Module): + # Set per instance in __init__; the default covers instances built without it. + _per_layer_dynamic_eplb_enabled = False + def __init__( self, *, From cbfccea10c09268b956804469c795d9f6c40d3c3 Mon Sep 17 00:00:00 2001 From: jintaop Date: Thu, 29 Oct 2026 13:23:54 -0700 Subject: [PATCH 56/59] [None][fix] Collect garbage before forking postprocess workers A long-lived MPI worker hosts one LLM after another. Cyclic garbage that the previous LLM left behind, CUDA objects included, could survive until the next LLM forked its postprocess workers; a child that collected it aborted with "CUDA error: initialization error", and test_llm_with_postprocess_parallel_and_result_handler then hung waiting for results nobody consumed. Collect in the parent before the fork. Signed-off-by: jintaop --- tensorrt_llm/executor/worker.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py index 602438cf13a7..9b1ef7f031e8 100644 --- a/tensorrt_llm/executor/worker.py +++ b/tensorrt_llm/executor/worker.py @@ -548,6 +548,10 @@ def notify_proxy_threads_to_quit(): [worker_queues.result_queue_addr]) assert result_queues is not None + # A reused worker process may still hold cyclic garbage from the LLM it + # hosted before, including CUDA objects. Collect it here: a forked + # postprocess worker that finalizes them aborts on CUDA. + gc.collect() postproc_worker_pool = ProcessPoolExecutor( max_workers=postproc_worker_config.num_postprocess_workers) for i in range(postproc_worker_config.num_postprocess_workers): From b7985115bc7dfd784ba3a7220f6593d5ef37ee1b Mon Sep 17 00:00:00 2001 From: jintaop Date: Thu, 29 Oct 2026 17:43:00 -0700 Subject: [PATCH 57/59] [None][fix] Keep ModelEngine.cleanup rank-local by default main's test_cleanup_releases_runner_model_before_release_gc expects cleanup() to be the finalizer's rank-local cleanup: release the model loader, then the graph captures without NCCL window owners, then the model references. Make that the default again and keep the collective order behind collective=True, which only the contract tests call; __del__ already passes collective=False. Signed-off-by: jintaop --- .../_torch/pyexecutor/model_engine.py | 30 +++++++++++-------- .../_torch/moe/test_dynamic_eplb_contracts.py | 8 +++-- 2 files changed, 23 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 7454ddca3fb0..b974140b43d2 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1125,11 +1125,7 @@ def finish_cleanup(self) -> None: if self._cleanup_done: return - model_loader = self.model_loader - if model_loader is not None: - model_loader.cleanup() - self.model_loader = None - + self._release_model_loader() self.moe_load_balancer = None # The runner, caller and scheduler retain the model, so clearing the # engine's attribute alone would leave the weights reachable. @@ -1143,19 +1139,29 @@ def finish_cleanup(self) -> None: release_gc() self._cleanup_done = True - def cleanup(self, *, collective: bool = True) -> None: - """Release model resources without rank-asymmetric collectives. + def _release_model_loader(self) -> None: + # Keep the handle live until its release succeeds, so a failed cleanup + # can be retried without releasing the loader twice. + model_loader = self.model_loader + if model_loader is not None: + model_loader.cleanup() + self.model_loader = None + + def cleanup(self, *, collective: bool = False) -> None: + """Release model resources. - Explicit terminal cleanup first releases local graph captures on every - rank, then passes that local result into the load-balancer's collective - preflight, and only then drops rank-local model state. Python finalizers - perform local cleanup only; worker-owned terminal cleanup is responsible - for collective resources. + By default this is the rank-local cleanup that Python finalizers run: + the model loader, then graph captures, then the model references. With + ``collective``, it first releases local graph captures on every rank, + then passes that local result into the load-balancer's collective + preflight, and only then drops rank-local model state; worker-owned + terminal cleanup is responsible for those collective resources. """ if self._cleanup_done: return if not collective: + self._release_model_loader() # Destruction timing is not rank-symmetric; only explicit shutdown # releases NCCL window owners. self.prepare_cleanup(release_nccl_window_owners=False) diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index 0111b2a3f9e7..b557dcbcb885 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -638,6 +638,7 @@ def test_model_engine_cleanup_shuts_manager_after_graphs_while_model_is_alive(): "prepare_cleanup", "shutdown_moe_load_balancer", "finish_cleanup", + "_release_model_loader", "cleanup", }, {"release_gc": lambda: events.append("release_gc")}, @@ -661,8 +662,8 @@ def shutdown(self, *, local_safe): engine.model = model engine.input_processor = object() - engine.cleanup() - engine.cleanup() + engine.cleanup(collective=True) + engine.cleanup(collective=True) assert events == [ "release_graphs", @@ -682,6 +683,7 @@ def test_model_engine_prepare_failure_participates_in_unsafe_preflight(): "prepare_cleanup", "shutdown_moe_load_balancer", "finish_cleanup", + "_release_model_loader", "cleanup", }, {"release_gc": lambda: events.append("release_gc")}, @@ -706,7 +708,7 @@ def shutdown(self, *, local_safe): engine.input_processor = object() with pytest.raises(RuntimeError, match="graph cleanup failed"): - engine.cleanup() + engine.cleanup(collective=True) assert events == [("manager.shutdown", False)] assert engine.moe_load_balancer is not None From 076a01afa21cf1b32e7b0494ce42b10d63ae8262 Mon Sep 17 00:00:00 2001 From: jintaop Date: Fri, 30 Oct 2026 01:17:06 -0700 Subject: [PATCH 58/59] [None][fix] Fit collective executor shutdown into the resource-rebind lifecycle main split PyExecutor teardown into release_resources(), rebind_resources() and a final shutdown() over a shared _release_execution_resources(). Carry the per-layer dynamic EPLB shutdown contract over to that structure: - shutdown() snapshots the model engines first, so the worker's collective terminal_cleanup() can still reach them after the executor drops its model_engine/draft_model_engine references. - A detected hang clears the snapshot so the worker never enters model collectives on a broken rank set. - A retried release stops the sampler's async worker once and does not repeat request publication, graph release or resource-manager shutdown. Both flags reset once a release completes, so the next release (after a rebind, or of resources a failed rebind bound) runs in full. The contract tests follow the new method names, and the stream-sharing test now checks every executor build of the engine, which is one since profiling and serving reuse a single executor. Signed-off-by: jintaop --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 104 ++++++++++++++++-- .../_torch/moe/test_dynamic_eplb_contracts.py | 16 +-- 2 files changed, 101 insertions(+), 19 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 28b3792ec224..c3adfddf9925 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2240,18 +2240,27 @@ def shutdown(self) -> None: """Finally close the executor; closed executors cannot be rebound.""" if self._closed: return + # Capture model owners before any rank-local shutdown step can fail; + # outer workers still need them for an all-rank unsafe preflight. + terminal_engines = list(getattr(self, "_terminal_model_engines", ())) + for engine in (self.model_engine, self.draft_model_engine): + if engine is not None and all(engine is not item + for item in terminal_engines): + terminal_engines.append(engine) + self._terminal_model_engines = tuple(terminal_engines) if not self._resources_released: if not self._release_execution_resources( release_model_owned_graphs=True): return else: - for engine in (self.model_engine, self.draft_model_engine): - if engine is not None: - engine.release_executor_resources( - release_model_owned_graphs=True) + for engine in self._terminal_model_engines: + engine.release_executor_resources( + release_model_owned_graphs=True) if torch.cuda.is_available(): torch.cuda.synchronize() self._closed = True + # The snapshot above keeps the model owners reachable for the outer + # worker's collective terminal cleanup. del self.model_engine if self.draft_model_engine is not None: del self.draft_model_engine @@ -2261,6 +2270,8 @@ def _release_execution_resources(self, release_model_owned_graphs: bool = False ) -> bool: """Stop execution and release GPU resources; retain model ownership.""" + if getattr(self, "_shutdown_runtime_complete", False): + return self._release_execution_pools() node_server = getattr(self, "_openengine_node_server", None) if node_server is not None: node_server.close() @@ -2294,6 +2305,7 @@ def _release_execution_resources(self, # enter model collectives on a broken rank set. self._pp_mpi_progress_stop.set() self._stop_mpi_progress_thread() + self._terminal_model_engines = () return False if self.worker_started: self.worker_thread.join() @@ -2320,7 +2332,7 @@ def _release_execution_resources(self, self.encoder_launch_executor = None # Stop the sampler's async worker, if it was used. A retried - # shutdown stops it only once. + # release stops it only once. if not getattr(self, "_shutdown_sampler_complete", False): if (isinstance(self.sampler, AsyncWorkerMixin) and self.sampler.async_worker_enabled()): @@ -2340,9 +2352,10 @@ def _release_execution_resources(self, # empty_cache), the subsequent CUDA graph teardown can trigger a # device-wide cudaErrorIllegalAddress when the driver touches metadata # for the now-freed memory regions. - for engine in self._terminal_model_engines: - engine.release_executor_resources( - release_model_owned_graphs=release_model_owned_graphs) + for engine in (self.model_engine, self.draft_model_engine): + if engine is not None: + engine.release_executor_resources( + release_model_owned_graphs=release_model_owned_graphs) # Ensure graph destruction has fully completed on device before # resource managers start freeing GPU-backed workspaces. if torch.cuda.is_available(): @@ -2353,16 +2366,83 @@ def _release_execution_resources(self, for manager in self.resource_manager.resource_managers.values(): if manager: manager.shutdown() + # Retry a later failure without repeating request publication, graph + # release, or resource-manager shutdown. + self._shutdown_runtime_complete = True + return self._release_execution_pools() + + def _release_execution_pools(self) -> bool: if self.virtual_memory_pools is not None: keys = list(self.virtual_memory_pools.keys()) for key in keys: - del virtual_memory_pools[key] - dwdp_manager = getattr(self, "dwdp_manager", None) - if dwdp_manager is not None: - dwdp_manager.__exit__(None, None, None) + del self.virtual_memory_pools[key] + if self.dwdp_manager is not None: + self.dwdp_manager.__exit__(None, None, None) self.dwdp_manager = None + # Released in full: the next release, after a rebind or of resources a + # failed rebind bound, starts over. + self._shutdown_runtime_complete = False + self._shutdown_sampler_complete = False return True + def terminal_cleanup(self, *, local_safe: bool = True) -> None: + """Collectively release model resources after all worker threads stop. + + Every rank first completes local graph preparation for every engine, + then participates in every manager preflight/collective in the same + order. Rank-local model and FC12 cache release happens only afterwards. + """ + engines = getattr(self, "_terminal_model_engines", ()) + if not engines: + return + + prepare_errors = [] + if local_safe: + for engine in engines: + try: + engine.prepare_cleanup() + except BaseException as error: + prepare_errors.append(error) + + collective_errors = [] + collective_safe = bool(local_safe and not prepare_errors) + for engine in engines: + try: + manager_safe = engine.shutdown_moe_load_balancer( + local_safe=collective_safe) + except BaseException as error: + collective_errors.append(error) + collective_safe = False + else: + collective_safe = collective_safe and manager_safe + + if prepare_errors: + raise prepare_errors[0] + if collective_errors: + raise collective_errors[0] + if not collective_safe: + raise RuntimeError( + "MoE load-balancer collective shutdown was not safe on every rank" + ) + + local_errors = [] + for engine in engines: + try: + engine.finish_cleanup() + except BaseException as error: + local_errors.append(error) + + # FC12 runners are rank-local and may be released only after every + # model-owned collective resource has been closed. + try: + self.release_rank_local_caches() + except BaseException as error: + local_errors.append(error) + + if local_errors: + raise local_errors[0] + self._terminal_model_engines = () + def local_cleanup(self) -> None: """Rank-local teardown for executors without collective EP resources. diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index b557dcbcb885..5396b4557412 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -878,7 +878,7 @@ def __exit__(self, *args): executor_cls = _extract_class( _TORCH_ROOT / "pyexecutor/py_executor.py", "PyExecutor", - {"shutdown", "_shutdown", "_finish_shutdown"}, + {"shutdown", "_release_execution_resources", "_release_execution_pools"}, { "AsyncWorkerMixin": AsyncWorkerMixin, "torch": SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: False)), @@ -895,6 +895,9 @@ def __exit__(self, *args): executor = executor_cls() executor.model_engine = engine executor.draft_model_engine = None + executor._closed = False + executor._resources_released = False + executor.kv_cache_transceiver = None executor.shutdown_all_ranks = True executor.can_shutdown = lambda: True executor.executor_request_queue = SimpleNamespace( @@ -923,7 +926,7 @@ def __exit__(self, *args): def terminal_cleanup(*, local_safe): events.append(("terminal_cleanup", local_safe)) if local_safe: - assert executor.model_engine is None + assert not hasattr(executor, "model_engine") assert executor._terminal_model_engines == (engine,) executor._terminal_model_engines = () @@ -961,10 +964,9 @@ def terminal_cleanup(*, local_safe): ("terminal_cleanup", True), ] assert executor.virtual_memory_pools == {} - assert executor.model_engine is None + assert not hasattr(executor, "model_engine") assert executor.draft_model_engine is None assert executor.dwdp_manager is None - assert executor._shutdown_sampler_complete assert executor._terminal_model_engines == () assert worker.engine is None @@ -1761,8 +1763,8 @@ def buffer(ptr): def test_executors_of_one_engine_share_the_execution_stream(): - # The planner binds the first MAIN stream for good, so the KV-cache - # estimation executor and the final executor must share one stream. + # The planner binds the first MAIN stream for good, so every executor + # built for the engine, across profiling and serving, uses its one stream. tree = ast.parse((_TORCH_ROOT / "pyexecutor/py_executor_creator.py").read_text()) create = next( node @@ -1776,7 +1778,7 @@ def test_executors_of_one_engine_share_the_execution_stream(): and isinstance(node.func, ast.Name) and node.func.id == "create_py_executor_instance" ] - assert len(calls) >= 2 + assert calls for call in calls: stream = next( (keyword.value for keyword in call.keywords if keyword.arg == "execution_stream"), From 23851dfdb72bfd0048078310f5c3d7e8b29777e0 Mon Sep 17 00:00:00 2001 From: jintaop Date: Fri, 30 Oct 2026 01:26:10 -0700 Subject: [PATCH 59/59] [None][fix] Address layer-wise benchmark review feedback - Per-layer dynamic EPLB records and replays expert IDs, so label them with num_experts instead of num_slots. With the load balancer off, slots are expert IDs as well, so one recording now drives both arms. - README: drop the num_slots replay restriction, and note that without replay the per-layer dynamic arm reuses one plan per layer, skips the helper copies after warmup and so times a best case. - Move the runner's load-balancer guards into _check_load_balancer_support and cover the accepted and rejected configurations with CPU tests. Signed-off-by: jintaop --- examples/layer_wise_benchmarks/README.md | 15 +++++--- .../_torch/moe/fused_moe/moe_scheduler.py | 19 +++++----- .../tools/layer_wise_benchmarks/runner.py | 35 ++++++++++--------- .../_torch/moe/test_dynamic_eplb_contracts.py | 6 ++++ .../test_layer_wise_benchmarks_runner_args.py | 27 ++++++++++++++ 5 files changed, 73 insertions(+), 29 deletions(-) diff --git a/examples/layer_wise_benchmarks/README.md b/examples/layer_wise_benchmarks/README.md index c9c70d311377..880dfb8e22b9 100644 --- a/examples/layer_wise_benchmarks/README.md +++ b/examples/layer_wise_benchmarks/README.md @@ -138,11 +138,16 @@ NP=8 ./mpi_launch.sh -x MEGAMOE_COMBINE_FORMAT=16e2m1xbf16 -x ENABLE_PERFECT_ROU # Per-layer dynamic EPLB runs without CUDA graphs only, and online standard EPLB is not # supported on a layer slice; the runner rejects both. # To replay the routing of an end-to-end run, add --replay-file-path (see "Performance -# alignment" below) to the first two arms only, each with a recording made with the same -# num_slots: a recording with the load balancer off (384 slots) cannot be replayed into the -# per-layer dynamic arm (408 slots), nor the reverse. Replay overrides the expert IDs that -# ENABLE_PERFECT_ROUTER=expert_ids substitutes, so run the third arm without it; run.py -# rejects that combination. Replay with MEGAMOE_CUTEDSL and the CTX phase is untested. +# alignment" below) to the first two arms only. Both arms record and replay expert IDs, so +# one recording, for example made with the load balancer off, drives both. Replay overrides +# the expert IDs that ENABLE_PERFECT_ROUTER=expert_ids substitutes, so run the third arm +# without it; run.py rejects that combination. Replay with MEGAMOE_CUTEDSL and the CTX +# phase is untested. +# Without replay every timed iteration reruns the same inputs, so the per-layer dynamic arm +# plans the same placement each time: after warmup its helper slots already hold the planned +# experts and the copies are skipped. That arm then times HALO-Q and the plan sync with +# almost no helper copies, a best case. Run this arm with replay to include the copies; +# compare the per-layer dynamic and balanced arms end to end. # Run with DeepEP A2A NP=4 ./mpi_launch.sh -x TRTLLM_FORCE_COMM_METHOD=DEEPEP ./run.sh config_ctx.yaml --moe-backend CUTEDSL diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py index f07aa1e87856..686322d1fff4 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py @@ -84,16 +84,17 @@ def _split_rows(t: torch.Tensor, sizes: List[int]) -> List[torch.Tensor]: def _record_routes( - moe: "ConfigurableMoE", routes: Optional[torch.Tensor] + moe: "ConfigurableMoE", routes: Optional[torch.Tensor], num_slots: int ) -> Optional[torch.Tensor]: - """Record ``routes`` for expert statistics and the layer-wise calibrator. + """Record ``routes``, IDs in ``[0, num_slots)``, for expert statistics and + the layer-wise calibrator. Returns ``routes``, or the routes the calibrator replays in their place. """ if routes is not None: ExpertStatistic.set_layer(moe.layer_idx) - ExpertStatistic.maybe_add_info(moe.num_slots, routes) - return get_calibrator().maybe_collect_or_replay_slots(moe.num_slots, routes) + ExpertStatistic.maybe_add_info(num_slots, routes) + return get_calibrator().maybe_collect_or_replay_slots(num_slots, routes) def _wait_input_ready(event: Optional[torch.cuda.Event]) -> None: @@ -1573,10 +1574,12 @@ def _forward_chunk( # ----- EPLB: update stats + remap expert ids -> slot ids ----- # Per-layer dynamic EPLB plans from the routes that the layer-wise - # calibrator records or replays, so they are recorded before the route; - # otherwise the recorded routes are the slots routed to. + # calibrator records or replays, so they are recorded before the route, + # as expert IDs; otherwise the recorded routes are the slots routed to. + # With the load balancer off, slots are expert IDs, so one recording + # drives both arms. if uses_replica_plan: - token_selected_experts = _record_routes(moe, token_selected_experts) + token_selected_experts = _record_routes(moe, token_selected_experts, moe.num_experts) token_selected_slots = moe._load_balancer_route(token_selected_experts, moe.use_dp) replica_plan = moe.layer_load_balancer.require_replica_plan() else: @@ -1594,7 +1597,7 @@ def _forward_chunk( token_selected_slots = moe._load_balancer_route(token_selected_experts, moe.use_dp) else: token_selected_slots = token_selected_experts - token_selected_slots = _record_routes(moe, token_selected_slots) + token_selected_slots = _record_routes(moe, token_selected_slots, moe.num_slots) replica_plan = None # ----- quantize / prepare ----- diff --git a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py index c48a96278a2a..78788cc45094 100644 --- a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py +++ b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py @@ -397,6 +397,23 @@ def forward_impl(*args, **kwargs): return forward_impl +def _check_load_balancer_support(load_balancer_config, use_cuda_graph: bool) -> None: + """Reject load balancers that cannot run on a layer slice driven by this harness.""" + load_balancer_mode = getattr(load_balancer_config, "mode", None) + if load_balancer_mode == "per_layer_dynamic" and use_cuda_graph: + raise ValueError( + "Per-layer dynamic EPLB does not support CUDA graphs: it binds one submission" + " stream and keeps host-side plan state that a graph replay cannot advance." + " Rerun with --no-use-cuda-graph." + ) + if load_balancer_mode == "standard" and load_balancer_config.layer_updates_per_iter > 0: + raise ValueError( + "Layer-wise benchmarks support per_layer_dynamic EPLB and static standard EPLB" + " (layer_updates_per_iter 0) only: online standard EPLB synchronizes every MoE" + " layer of the model, but only a slice of the layers runs here." + ) + + class Runner: def __init__( self, @@ -451,22 +468,8 @@ def __init__( dtype=kv_cache_dtype, mamba_ssm_cache_dtype=mamba_ssm_cache_dtype ), ) - # Reject load balancers that cannot run on a layer slice driven by this - # harness. The validator above has resolved a YAML path to its config. - load_balancer_config = llm_args.moe_config.load_balancer - load_balancer_mode = getattr(load_balancer_config, "mode", None) - if load_balancer_mode == "per_layer_dynamic" and use_cuda_graph: - raise ValueError( - "Per-layer dynamic EPLB does not support CUDA graphs: it binds one submission" - " stream and keeps host-side plan state that a graph replay cannot advance." - " Rerun with --no-use-cuda-graph." - ) - if load_balancer_mode == "standard" and load_balancer_config.layer_updates_per_iter > 0: - raise ValueError( - "Layer-wise benchmarks support per_layer_dynamic EPLB and static standard EPLB" - " (layer_updates_per_iter 0) only: online standard EPLB synchronizes every MoE" - " layer of the model, but only a slice of the layers runs here." - ) + # The validator above has resolved a YAML path to its config. + _check_load_balancer_support(llm_args.moe_config.load_balancer, use_cuda_graph) with ( self.scaled_from_ctx(scaled_from, mapping), self.vision_config_ctx(vision_config), diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py index 5396b4557412..1ecd11a30fd0 100644 --- a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -2379,6 +2379,7 @@ def __init__(self, name): def test_fused_comm_chunk_routes_once_after_recording_the_planner_input(mode): source = _TORCH_ROOT / "moe/fused_moe/moe_scheduler.py" trace = [] + labels = [] class ExpertStatistic: @staticmethod @@ -2387,9 +2388,11 @@ def set_layer(layer_idx): @staticmethod def maybe_add_info(num_slots, routes): + labels.append(num_slots) trace.append(("statistic", routes.name)) def maybe_collect_or_replay_slots(num_slots, routes): + labels.append(num_slots) trace.append(("calibrator", routes.name)) return _ChunkRoutes("replayed " + routes.name) @@ -2429,6 +2432,7 @@ def run_moe(ctx): apply_router_weight_on_input=False, mapping=SimpleNamespace(moe_ep_rank=0), layer_idx=7, + num_experts=16, num_slots=24, use_dp=True, layer_load_balancer=layer, @@ -2506,6 +2510,8 @@ def run_moe(ctx): "done_set_cpu", ] assert trace == expected + # Routes recorded before the route are expert IDs; after it, slots. + assert labels == ([16, 16] if mode == "per_layer_dynamic" else [24, 24]) def test_base_worker_finalizes_rank_follower_pyexecutor(): diff --git a/tests/unittest/tools/test_layer_wise_benchmarks_runner_args.py b/tests/unittest/tools/test_layer_wise_benchmarks_runner_args.py index f13b40e280c2..63faf004f263 100644 --- a/tests/unittest/tools/test_layer_wise_benchmarks_runner_args.py +++ b/tests/unittest/tools/test_layer_wise_benchmarks_runner_args.py @@ -20,6 +20,7 @@ import pytest +from tensorrt_llm.llmapi.llm_args import MoeLoadBalancerConfig, MoePerLayerDynamicLoadBalancerConfig from tensorrt_llm.mapping import Mapping from tensorrt_llm.tools.layer_wise_benchmarks import runner as runner_module @@ -87,3 +88,29 @@ def __init__(self, **kwargs): assert llm_args.max_beam_width == 3 # ... and are resolved from the text-only config the model loads with. assert recorded["config_loader"] == "load_pretrained_config_text_only" + + +@pytest.mark.parametrize( + ("load_balancer", "use_cuda_graph", "error"), + [ + (MoePerLayerDynamicLoadBalancerConfig(num_slots=408), True, "does not support CUDA graphs"), + (MoeLoadBalancerConfig(num_slots=384, layer_updates_per_iter=1), False, "static standard"), + ], +) +def test_runner_rejects_load_balancers_a_layer_slice_cannot_run( + load_balancer, use_cuda_graph, error +): + with pytest.raises(ValueError, match=error): + runner_module._check_load_balancer_support(load_balancer, use_cuda_graph) + + +@pytest.mark.parametrize( + ("load_balancer", "use_cuda_graph"), + [ + (None, True), + (MoePerLayerDynamicLoadBalancerConfig(num_slots=408), False), + (MoeLoadBalancerConfig(num_slots=384, layer_updates_per_iter=0), True), + ], +) +def test_runner_accepts_load_balancers_a_layer_slice_can_run(load_balancer, use_cuda_graph): + runner_module._check_load_balancer_support(load_balancer, use_cuda_graph)