From c980c4fb7a69768df843f758378ff980567b7cc1 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Sun, 1 Nov 2026 10:03:50 -0800 Subject: [PATCH] feat: extend Rubin GLM low-latency targets within the GLM family Signed-off-by: Dongfeng Yu --- .github/CODEOWNERS | 8 +- .../kernels/glmMoeDsa/glmLowMBf16Moe.cu | 838 --------- .../kernels/glmMoeDsa/glmMoeFusedRouter.cu | 574 ------ .../kernels/glmMoeDsa/glmRouterGemm.cu | 244 --- .../glmDenseMlpSmallM.cu | 2 +- .../glmDenseMlpSmallM.h | 0 .../glmDraftSampleGlue.cu | 366 +++- .../glmDraftSampleGlue.h | 0 .../glmDsaHistTopK.cu | 2 +- .../glmDsaHistTopK.h | 0 .../glmFusedAGemm.cu | 2 +- .../glmFusedAGemm.h | 0 .../glmFusedEmbeddingLookup.cu | 2 +- .../glmFusedEmbeddingLookup.h | 0 .../glmMoeDsaRubin/glmIntraNodeLamport.cuh | 212 +++ .../glmMoeDsaRubin/glmIntraNodeReduce.cu | 600 ++++++ .../glmMoeDsaRubin/glmIntraNodeReduce.h | 53 + .../glmLamport.cuh | 0 .../kernels/glmMoeDsaRubin/glmLowMBf16Moe.cu | 1664 +++++++++++++++++ .../glmLowMBf16Moe.h | 34 + .../glmMlaKvHalfAppend.cu | 2 +- .../glmMlaKvHalfAppend.h | 0 .../glmMnnvlAllgather.cu | 4 +- .../glmMnnvlAllgather.h | 0 .../glmMnnvlHop.cu | 4 +- .../glmMnnvlHop.h | 0 .../glmMoeDsaRubin/glmMoeFusedRouter.cu | 1415 ++++++++++++++ .../glmMoeFusedRouter.h | 11 +- .../glmMoeGemm2K128.cu | 2 +- .../glmMoeGemm2K128.h | 0 .../kernels/glmMoeDsaRubin/glmMoeGemm2K256.cu | 708 +++++++ .../kernels/glmMoeDsaRubin/glmMoeGemm2K256.h | 67 + .../kernels/glmMoeDsaRubin/glmRouterGemm.cu | 536 ++++++ .../glmRouterGemm.h | 10 +- .../glmSparseMlaFlashDecode.cu | 667 ++++++- .../glmSparseMlaFlashDecode.h | 32 +- .../glmVerifyHeadAccept.cu | 2 +- .../glmVerifyHeadAccept.h | 0 .../glmMoeDsaRubin/glmVerifyHeadSample.cu | 1204 ++++++++++++ .../glmMoeDsaRubin/glmVerifyHeadSample.h | 52 + cpp/tensorrt_llm/thop/CMakeLists.txt | 2 +- .../CMakeLists.txt | 8 +- .../glmDenseMlpSmallMOp.cpp | 2 +- .../glmDraftSampleGlueOp.cpp | 2 +- .../glmDsaHistTopKOp.cpp | 2 +- .../glmFusedAGemmOp.cpp | 2 +- .../glmFusedEmbeddingOp.cpp | 2 +- .../glmMoeDsaRubin/glmIntraNodeReduceOp.cpp | 217 +++ .../glmLowMBf16MoeOp.cpp | 60 +- .../glmMlaKvHalfAppendOp.cpp | 2 +- .../glmMnnvlAllgatherOp.cpp | 4 +- .../glmMnnvlHopOp.cpp | 10 +- .../glmRoutedExpertsOp.cpp | 45 +- .../glmRouterGemmOp.cpp | 6 +- .../glmSparseMlaFlashDecodeOp.cpp | 132 +- .../glmVerifyHeadAcceptOp.cpp | 2 +- .../glmMoeDsaRubin/glmVerifyHeadSampleOp.cpp | 82 + .../__init__.py | 0 .../_full_context.py | 210 +++ .../_model.py | 0 .../_mtp.py | 18 + .../_mtp_worker.py | 568 ++++++ .../_small_batch_draft.py | 475 +++++ .../_speculative.py | 0 .../_weight_loader.py | 0 .../modeling.py | 830 ++++---- .../ops/__init__.py | 0 .../ops/_cpp_fakes.py | 98 + .../ops/attention/__init__.py | 0 .../ops/attention/mla_kv_half_append.py | 0 .../ops/attention/sparse_mla_flash_decode.py | 233 ++- .../ops/comm/__init__.py | 0 .../ops/comm/_ipc_workspace.py | 176 ++ .../ops/comm/_mnnvl_workspace.py | 0 .../ops/comm/intra_node_moe_finalize.py | 110 ++ .../ops/comm/intra_node_reduce.py | 182 ++ .../ops/comm/mnnvl_allgather.py | 0 .../ops/comm/mnnvl_hop.py | 0 .../ops/comm/mnnvl_moe_finalize_hop.py | 0 .../ops/comm/pre_moe_fp4_linear.py | 72 + .../ops/dsa/__init__.py | 0 .../ops/dsa/_fp8_paged_mqa_logits.py | 0 .../ops/dsa/_fp8_paged_mqa_logits_kernel.py | 0 .../ops/dsa/_full_context_topk.py | 154 ++ .../ops/dsa/_hist_topk.py | 2 +- .../ops/dsa/_indexer_prologue.py | 0 .../ops/dsa/_indexer_reference.py | 0 .../ops/dsa/_mtp_topk_stash.py | 0 .../ops/dsa/convert_req_index_to_global.py | 0 .../ops/dsa/fp8_paged_mqa_logits.py | 14 +- .../ops/dsa/hist_shard_meta.py | 0 .../ops/dsa/hist_topk.py | 0 .../ops/dsa/indexer_k_tail.py | 0 .../ops/dsa/indexer_prologue.py | 0 .../ops/dsa/paged_mqa_logits_schedule.py | 0 .../ops/dsa/stash_last_accepted_rows.py | 10 +- .../ops/embedding/__init__.py | 0 .../ops/embedding/fused_embedding.py | 0 .../ops/gemm/__init__.py | 0 .../ops/gemm/_low_m_splitk.py | 182 +- .../ops/gemm/_low_m_splitk_fused.py | 0 .../ops/gemm/dense_mlp_small_m.py | 0 .../ops/gemm/draft_hidden_half.py | 48 +- .../ops/gemm/fused_a_gemm.py | 47 +- .../ops/gemm/indexer_q_tail.py | 0 .../ops/gemm/low_m_gemm.py | 41 +- .../ops/gemm/o_proj_publish.py | 0 .../ops/gemm/q_b_fp8_rope.py | 8 + .../ops/gemm/router_gemm.py | 21 +- .../ops/moe/__init__.py | 0 .../ops/moe/lowm_bf16_moe.py | 66 +- .../ops/moe/routed_experts.py | 98 +- .../ops/norm/__init__.py | 0 .../ops/quantization/__init__.py | 0 .../ops/quantization/fp4_quantize.py | 0 .../ops/spec/__init__.py | 0 .../ops/spec/_draft_exchange.py | 0 .../ops/spec/_draft_glue_rmsnorm.py | 0 .../ops/spec/_draft_step_state.py | 0 .../ops/spec/_draft_token_table.py | 0 .../ops/spec/_gather_last_accepted_rows.py | 0 .../ops/spec/draft_ehproj_table.py | 0 .../ops/spec/draft_glue_rmsnorm_cat.py | 0 .../ops/spec/draft_sample_glue.py | 30 +- .../ops/spec/gather_last_accepted_rows.py | 0 .../ops/spec/verify_head_accept.py | 0 .../ops/spec/verify_head_sample.py | 178 ++ .../weights.py | 2 +- .../modeling_v2/models/glm_moe_dsa/routing.py | 22 +- tests/integration/defs/.test_durations | 4 +- .../accuracy/modeling_v2/test_glm_moe_dsa.py | 785 +++++--- .../integration/test_lists/test-db/l0_cpu.yml | 2 +- .../test_lists/test-db/l0_gr100.yml | 14 +- tests/integration/test_lists/waives.txt | 3 - .../parallel_hw_agnostic/test_custom_ops.py | 2 +- .../test_modeling_v2_mla_kv_half_append.py | 4 +- ...est_modeling_v2_sparse_mla_flash_decode.py | 235 ++- .../comm/_intra_node_reduce_op_matrix.py | 349 ++++ .../comm/_mnnvl_allgather_op_matrix.py | 2 +- .../comm/_mnnvl_hop_op_matrix.py | 153 +- .../comm/_mnnvl_moe_finalize_hop_op_matrix.py | 111 +- .../comm/_mnnvl_ranks.py | 59 + .../comm/_rank_job.py | 0 ...modeling_v2_intra_node_reduce_op_matrix.py | 21 + ...t_modeling_v2_mnnvl_allgather_op_matrix.py | 0 .../test_modeling_v2_mnnvl_hop_op_matrix.py | 0 ...ing_v2_mnnvl_moe_finalize_hop_op_matrix.py | 0 .../dsa/_glm_rank_job.py | 0 .../dsa/_hist_topk_op_matrix.py | 4 +- .../dsa/_indexer_inputs.py | 0 ...modeling_v2_convert_req_index_to_global.py | 2 +- .../test_modeling_v2_fp8_paged_mqa_logits.py | 4 +- .../dsa/test_modeling_v2_full_context_topk.py | 163 ++ .../test_modeling_v2_hist_buffer_lifecycle.py | 8 +- .../dsa/test_modeling_v2_hist_shard_meta.py | 4 +- .../test_modeling_v2_hist_topk_op_matrix.py | 3 +- .../dsa/test_modeling_v2_indexer_k_tail.py | 4 +- .../dsa/test_modeling_v2_indexer_prologue.py | 6 +- ...t_modeling_v2_paged_mqa_logits_schedule.py | 2 +- ...st_modeling_v2_stash_last_accepted_rows.py | 40 +- .../test_modeling_v2_fused_embedding.py | 4 +- .../test_modeling_v2_dense_mlp_small_m.py | 4 +- .../test_modeling_v2_draft_hidden_half.py | 69 +- .../gemm/test_modeling_v2_fused_a_gemm.py | 14 +- .../gemm/test_modeling_v2_indexer_q_tail.py | 2 +- .../gemm/test_modeling_v2_low_m_gemm.py | 85 +- .../gemm/test_modeling_v2_o_proj_publish.py | 4 +- .../gemm/test_modeling_v2_q_b_fp8_rope.py | 2 +- .../gemm/test_modeling_v2_router_gemm.py | 32 +- .../moe/test_modeling_v2_lowm_bf16_moe.py | 44 +- .../moe/test_modeling_v2_routed_experts.py | 187 +- .../test_modeling_v2_glm_fp4_quantize.py | 2 +- .../spec/_draft_step_op_matrix.py | 342 +++- .../spec/_draft_step_rank_job.py | 0 .../test_modeling_v2_draft_ehproj_table.py | 6 +- .../spec/test_modeling_v2_draft_exchange.py | 2 +- ...test_modeling_v2_draft_glue_rmsnorm_cat.py | 2 +- .../test_modeling_v2_draft_sample_glue.py | 6 +- .../test_modeling_v2_draft_token_table.py | 4 +- ...t_modeling_v2_gather_last_accepted_rows.py | 2 +- .../spec/test_modeling_v2_small_batch_mtp.py | 605 ++++++ .../test_modeling_v2_verify_head_accept.py | 2 +- .../test_modeling_v2_verify_head_sample.py | 270 +++ .../models/glm_moe_dsa/test_routing.py | 19 +- .../modeling_v2/test_modeling_v2_claims.py | 2 +- 185 files changed, 14563 insertions(+), 2924 deletions(-) delete mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.cu delete mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmMoeFusedRouter.cu delete mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.cu rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmDenseMlpSmallM.cu (99%) rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmDenseMlpSmallM.h (100%) rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmDraftSampleGlue.cu (71%) rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmDraftSampleGlue.h (100%) rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmDsaHistTopK.cu (99%) rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmDsaHistTopK.h (100%) rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmFusedAGemm.cu (99%) rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmFusedAGemm.h (100%) rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmFusedEmbeddingLookup.cu (97%) rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmFusedEmbeddingLookup.h (100%) create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmIntraNodeLamport.cuh create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmIntraNodeReduce.cu create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmIntraNodeReduce.h rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmLamport.cuh (100%) create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmLowMBf16Moe.cu rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmLowMBf16Moe.h (61%) rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmMlaKvHalfAppend.cu (99%) rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmMlaKvHalfAppend.h (100%) rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmMnnvlAllgather.cu (98%) rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmMnnvlAllgather.h (100%) rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmMnnvlHop.cu (99%) rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmMnnvlHop.h (100%) create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeFusedRouter.cu rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmMoeFusedRouter.h (81%) rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmMoeGemm2K128.cu (99%) rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmMoeGemm2K128.h (100%) create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeGemm2K256.cu create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeGemm2K256.h create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmRouterGemm.cu rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmRouterGemm.h (83%) rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmSparseMlaFlashDecode.cu (54%) rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmSparseMlaFlashDecode.h (72%) rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmVerifyHeadAccept.cu (98%) rename cpp/tensorrt_llm/kernels/{glmMoeDsa => glmMoeDsaRubin}/glmVerifyHeadAccept.h (100%) create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmVerifyHeadSample.cu create mode 100644 cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmVerifyHeadSample.h rename cpp/tensorrt_llm/thop/{glmMoeDsa => glmMoeDsaRubin}/CMakeLists.txt (76%) rename cpp/tensorrt_llm/thop/{glmMoeDsa => glmMoeDsaRubin}/glmDenseMlpSmallMOp.cpp (98%) rename cpp/tensorrt_llm/thop/{glmMoeDsa => glmMoeDsaRubin}/glmDraftSampleGlueOp.cpp (99%) rename cpp/tensorrt_llm/thop/{glmMoeDsa => glmMoeDsaRubin}/glmDsaHistTopKOp.cpp (99%) rename cpp/tensorrt_llm/thop/{glmMoeDsa => glmMoeDsaRubin}/glmFusedAGemmOp.cpp (98%) rename cpp/tensorrt_llm/thop/{glmMoeDsa => glmMoeDsaRubin}/glmFusedEmbeddingOp.cpp (97%) create mode 100644 cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmIntraNodeReduceOp.cpp rename cpp/tensorrt_llm/thop/{glmMoeDsa => glmMoeDsaRubin}/glmLowMBf16MoeOp.cpp (61%) rename cpp/tensorrt_llm/thop/{glmMoeDsa => glmMoeDsaRubin}/glmMlaKvHalfAppendOp.cpp (99%) rename cpp/tensorrt_llm/thop/{glmMoeDsa => glmMoeDsaRubin}/glmMnnvlAllgatherOp.cpp (97%) rename cpp/tensorrt_llm/thop/{glmMoeDsa => glmMoeDsaRubin}/glmMnnvlHopOp.cpp (96%) rename cpp/tensorrt_llm/thop/{glmMoeDsa => glmMoeDsaRubin}/glmRoutedExpertsOp.cpp (94%) rename cpp/tensorrt_llm/thop/{glmMoeDsa => glmMoeDsaRubin}/glmRouterGemmOp.cpp (96%) rename cpp/tensorrt_llm/thop/{glmMoeDsa => glmMoeDsaRubin}/glmSparseMlaFlashDecodeOp.cpp (51%) rename cpp/tensorrt_llm/thop/{glmMoeDsa => glmMoeDsaRubin}/glmVerifyHeadAcceptOp.cpp (98%) create mode 100644 cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmVerifyHeadSampleOp.cpp rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/__init__.py (100%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/_full_context.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/_model.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/_mtp.py (95%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/_mtp_worker.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/_small_batch_draft.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/_speculative.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/_weight_loader.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/modeling.py (75%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/_cpp_fakes.py (64%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/attention/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/attention/mla_kv_half_append.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/attention/sparse_mla_flash_decode.py (64%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/comm/__init__.py (100%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/_ipc_workspace.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/comm/_mnnvl_workspace.py (100%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/intra_node_moe_finalize.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/intra_node_reduce.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/comm/mnnvl_allgather.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/comm/mnnvl_hop.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/comm/mnnvl_moe_finalize_hop.py (100%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/pre_moe_fp4_linear.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/dsa/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/dsa/_fp8_paged_mqa_logits.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/dsa/_fp8_paged_mqa_logits_kernel.py (100%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/_full_context_topk.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/dsa/_hist_topk.py (99%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/dsa/_indexer_prologue.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/dsa/_indexer_reference.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/dsa/_mtp_topk_stash.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/dsa/convert_req_index_to_global.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/dsa/fp8_paged_mqa_logits.py (94%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/dsa/hist_shard_meta.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/dsa/hist_topk.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/dsa/indexer_k_tail.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/dsa/indexer_prologue.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/dsa/paged_mqa_logits_schedule.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/dsa/stash_last_accepted_rows.py (89%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/embedding/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/embedding/fused_embedding.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/gemm/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/gemm/_low_m_splitk.py (90%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/gemm/_low_m_splitk_fused.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/gemm/dense_mlp_small_m.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/gemm/draft_hidden_half.py (74%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/gemm/fused_a_gemm.py (79%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/gemm/indexer_q_tail.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/gemm/low_m_gemm.py (77%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/gemm/o_proj_publish.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/gemm/q_b_fp8_rope.py (96%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/gemm/router_gemm.py (84%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/moe/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/moe/lowm_bf16_moe.py (65%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/moe/routed_experts.py (89%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/norm/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/quantization/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/quantization/fp4_quantize.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/spec/__init__.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/spec/_draft_exchange.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/spec/_draft_glue_rmsnorm.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/spec/_draft_step_state.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/spec/_draft_token_table.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/spec/_gather_last_accepted_rows.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/spec/draft_ehproj_table.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/spec/draft_glue_rmsnorm_cat.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/spec/draft_sample_glue.py (92%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/spec/gather_last_accepted_rows.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/ops/spec/verify_head_accept.py (100%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/verify_head_sample.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/weights.py (98%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/attention/test_modeling_v2_mla_kv_half_append.py (98%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/attention/test_modeling_v2_sparse_mla_flash_decode.py (72%) create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/_intra_node_reduce_op_matrix.py rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/comm/_mnnvl_allgather_op_matrix.py (99%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/comm/_mnnvl_hop_op_matrix.py (76%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/comm/_mnnvl_moe_finalize_hop_op_matrix.py (79%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/comm/_mnnvl_ranks.py (76%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/comm/_rank_job.py (100%) create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/test_modeling_v2_intra_node_reduce_op_matrix.py rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/comm/test_modeling_v2_mnnvl_allgather_op_matrix.py (100%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/comm/test_modeling_v2_mnnvl_hop_op_matrix.py (100%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/comm/test_modeling_v2_mnnvl_moe_finalize_hop_op_matrix.py (100%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/dsa/_glm_rank_job.py (100%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/dsa/_hist_topk_op_matrix.py (99%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/dsa/_indexer_inputs.py (100%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/dsa/test_modeling_v2_convert_req_index_to_global.py (97%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/dsa/test_modeling_v2_fp8_paged_mqa_logits.py (97%) create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_full_context_topk.py rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/dsa/test_modeling_v2_hist_buffer_lifecycle.py (94%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/dsa/test_modeling_v2_hist_shard_meta.py (95%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/dsa/test_modeling_v2_hist_topk_op_matrix.py (86%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/dsa/test_modeling_v2_indexer_k_tail.py (98%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/dsa/test_modeling_v2_indexer_prologue.py (96%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/dsa/test_modeling_v2_paged_mqa_logits_schedule.py (96%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/dsa/test_modeling_v2_stash_last_accepted_rows.py (62%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/embedding/test_modeling_v2_fused_embedding.py (97%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/gemm/test_modeling_v2_dense_mlp_small_m.py (98%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/gemm/test_modeling_v2_draft_hidden_half.py (69%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/gemm/test_modeling_v2_fused_a_gemm.py (95%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/gemm/test_modeling_v2_indexer_q_tail.py (98%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/gemm/test_modeling_v2_low_m_gemm.py (57%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/gemm/test_modeling_v2_o_proj_publish.py (97%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/gemm/test_modeling_v2_q_b_fp8_rope.py (98%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/gemm/test_modeling_v2_router_gemm.py (87%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/moe/test_modeling_v2_lowm_bf16_moe.py (79%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/moe/test_modeling_v2_routed_experts.py (81%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/quantization/test_modeling_v2_glm_fp4_quantize.py (96%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/spec/_draft_step_op_matrix.py (65%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/spec/_draft_step_rank_job.py (100%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/spec/test_modeling_v2_draft_ehproj_table.py (80%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/spec/test_modeling_v2_draft_exchange.py (98%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py (98%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/spec/test_modeling_v2_draft_sample_glue.py (80%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/spec/test_modeling_v2_draft_token_table.py (95%) rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/spec/test_modeling_v2_gather_last_accepted_rows.py (97%) create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_small_batch_mtp.py rename tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/{glm_5_2_nvfp4__sm_107__tp16 => glm_5_2_nvfp4__sm_107__low_latency}/spec/test_modeling_v2_verify_head_accept.py (98%) create mode 100644 tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_verify_head_sample.py diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 4ca9dc19d450..92c370ab5ea0 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -522,10 +522,10 @@ /tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tep4 @erictsai-nv # GLM SM107 target and its native kernels/bindings -/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16 @dongfengy @Tabrizian @xwang233 -/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16 @dongfengy @Tabrizian @xwang233 -/cpp/tensorrt_llm/kernels/glmMoeDsa @dongfengy @Tabrizian @xwang233 -/cpp/tensorrt_llm/thop/glmMoeDsa/ @dongfengy @Tabrizian @xwang233 +/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency @dongfengy @Tabrizian @xwang233 +/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency @dongfengy @Tabrizian @xwang233 +/cpp/tensorrt_llm/kernels/glmMoeDsaRubin @dongfengy @Tabrizian @xwang233 +/cpp/tensorrt_llm/thop/glmMoeDsaRubin/ @dongfengy @Tabrizian @xwang233 # Uncomment only on release branches (e.g. release/0.19): # * @NVIDIA/trt-llm-release-branch-approval diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.cu deleted file mode 100644 index bc67707178a8..000000000000 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.cu +++ /dev/null @@ -1,838 +0,0 @@ -/* - * 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/glmMoeDsa/glmLowMBf16Moe.h" - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/envUtils.h" -#include "tensorrt_llm/kernels/moeTopKFuncs.cuh" - -#include -#include -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels::glm_moe_dsa -{ -namespace -{ -namespace cg = cooperative_groups; - -using Bf16 = __nv_bfloat16; - -constexpr int kNumExperts = kLowMBf16MoeNumExperts; -constexpr int kTopK = kLowMBf16MoeTopK; -static_assert(kNumExperts % 32 == 0, "routing gives each lane of one warp kNumExperts / 32 candidates"); - -// FC1: 8 consumer warps and one producer warp per CTA. Consumer thread t reads the 16-byte pieces t, -// t + kFc1ConsumerThreads, ... of x and of every weight row, so the hidden size is a multiple of kHiddenGranule and at -// most kMaxPiecesPerThread of them. -constexpr int kFc1ConsumerWarps = 8; -constexpr int kFc1ConsumerThreads = kFc1ConsumerWarps * 32; -constexpr int kFc1Threads = kFc1ConsumerThreads + 32; -constexpr int kPieceElems = 8; -constexpr int kHiddenGranule = kFc1ConsumerThreads * kPieceElems; -constexpr int kMaxPiecesPerThread = 4; -// The FC1 ring holds whole weight rows, as many as fit in kFc1RingBytes and at most kMaxStages. At hidden size 6144 -// and interSize 256 that is 7 rows of 12,288 B, and an FC1 CTA (ring, row sums, barriers and the 1,024 B the runtime -// reserves per CTA) takes 87,840 B, so one FC2 CTA (138,312 B) fits beside it in the 233,472 B of an SM. -constexpr int kFc1RingBytes = 7 * 6144 * static_cast(sizeof(Bf16)); -constexpr int kMaxStages = 8; - -// FC2: a warp per routed expert. -constexpr int kFc2Threads = 256; -constexpr int kFc2Warps = kFc2Threads / 32; -static_assert(kTopK == kFc2Warps, "the FC2 kernel gives each routed expert one warp"); - -// FC2: mma.sync m16n8k16 tiles have 16 matrix rows; one 32-element K step is two mma calls (see mmaStep). A warp per -// routed expert walks the intermediate dimension in pairs of K steps, one per accumulator. -constexpr int kMmaRows = 16; -constexpr int kStepK = 32; -constexpr int kFc2StepPairK = 2 * kStepK; -// The largest intermediate size per rank whose FC2 working set fits in shared memory (kTiles = 1). -constexpr int kMaxInterSize = 512; - -// The routed scaling factor as the FC2 kernel takes it: the double value of noaux_tc_op's FP64 weight expression, its -// float value, and whether routingWeight's FP32 path applies to it (decided on the host, see makeRoutingScale). -struct RoutingScale -{ - double value; - float valueFloat; - int fastPath; -}; - -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) -constexpr int kCandidates = kNumExperts / 32; -constexpr int kFc1ProducerWarp = kFc1ConsumerWarps; -// Named barrier at which the FC1 consumer warps wait while the producer selects the experts and issues the first rows. -constexpr int kFc1ParkBarrier = 1; -// The FC2 kernel holds its routing computation until this long after its entry; see lowMBf16MoeFc2Kernel. -constexpr uint64_t kFc2RouteHoldNs = 1000; - -using TopKKey = reduce_topk::TopKRedType; - -// Same expression as sigmoid_accurate in noAuxTcKernels.cu. -__device__ __forceinline__ float sigmoidAccurate(float x) -{ - return 0.5f * tanhf(0.5f * x) + 0.5f; -} - -__device__ __forceinline__ uint32_t warpMax(uint32_t value) -{ - uint32_t result; - asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n" : "=r"(result) : "r"(value)); - return result; -} - -// The correction bias of the experts a lane holds (i * 32 + lane) as float. A constant, so the kernels may read it -// before their grid-dependency wait. -template -__device__ __forceinline__ void loadBias(BiasT const* bias, int lane, float (&values)[kCandidates]) -{ -#pragma unroll - for (int i = 0; i < kCandidates; ++i) - { - values[i] = static_cast(bias[i * 32 + lane]); - } -} - -// The routing keys of noaux_tc_op for n_group == 1 (deepseek_v3_topk_kernel, single group) of the experts a lane -// holds: the sigmoid score plus the correction bias, with the expressions of that kernel, packed with the expert index -// as reduceTopK (moeTopKFuncs.cuh) packs them. Stores the sigmoid scores when sigmoidScores is not null. The logits -// are read through L2. -__device__ __forceinline__ void buildKeys(float const* routerLogits, float const (&biasValues)[kCandidates], int lane, - TopKKey (&keys)[kCandidates], float* sigmoidScores) -{ -#pragma unroll - for (int i = 0; i < kCandidates; ++i) - { - int const expert = i * 32 + lane; - float const sigmoid = sigmoidAccurate(__ldcg(routerLogits + expert)); - if (sigmoidScores != nullptr) - { - sigmoidScores[expert] = sigmoid; - } - keys[i] = TopKKey{sigmoid + biasValues[i], expert}; - } -} - -// The top kTopK of one full warp's keys in the order of reduceTopK: descending score, on equal scores the lower index -// first. The per-lane sort is reduceTopK's; each round then takes the warp maximum of the lanes' leading keys with two -// 32-bit reductions, of the score half and then of the index half among the lanes holding the maximum score, which is -// the maximum of reduceTopK's 64-bit reduction. So the experts and their order are noaux_tc_op's. Returns the k-th -// expert on lane k < kTopK, 0 on the other lanes. -__device__ __forceinline__ int32_t selectTopK(TopKKey (&keys)[kCandidates], int lane) -{ - reduce_topk::Sort::run(keys); - - int32_t laneExpert = 0; - uint64_t packedMax = 0; -#pragma unroll - for (int k = 0; k < kTopK; ++k) - { - if (k > 0 && packedMax == keys[0].compValIdx) - { - // This lane held the previous maximum: drop it, as reduceTopK does. -#pragma unroll - for (int i = 0; i + 1 < kCandidates; ++i) - { - keys[i] = keys[i + 1]; - } - keys[kCandidates - 1] = TopKKey{-INFINITY, (kCandidates - 1) * 32 + lane}; - } - uint32_t const scoreBits = static_cast(keys[0].compValIdx >> 32); - uint32_t const maxScoreBits = warpMax(scoreBits); - uint32_t const indexBits = warpMax(scoreBits == maxScoreBits ? static_cast(keys[0].compValIdx) : 0U); - packedMax = (static_cast(maxScoreBits) << 32) | indexBits; - if (lane == k) - { - laneExpert = TopKKey::kMaxIdx - static_cast(indexBits & 0xFFFF); - } - } - return laneExpert; -} - -// noaux_tc_op's routing weight, static_cast(s * c / (S + 1e-20)) in double with s the expert's sigmoid score, -// S the sum of the selected scores and c the routed scaling factor, without FP64 instructions on the common path: FP64 -// in the FC2 CTAs drops the SM clock for the whole FC1 window (measured on Rubin: 2.38 -> 1.6 GHz). When c is a float -// with at most 20 significant bits (fastPath) and S >= 2^-13, the double expression is the float rounding of the exact -// quotient s * c / S: s * c is exact in double, S + 1e-20 rounds to S, and rounding to double first cannot move a -// quotient of these operands across a float rounding boundary except at an exact midpoint, which rounds the same -// either way. The quotient here comes from an exact two-float product, a correctly rounded division and its exact -// remainder; it is returned when it lies more than 2^-20 ulp inside its rounding interval (and is not a power of two, -// where the interval is lopsided). Otherwise, and for scores below 2^-60 (where the FP32 intermediates could lose -// bits), the FP64 expression runs, out of line so that none of it is scheduled on the common path. -__device__ __noinline__ float routingWeightFp64(float s, float S, double c) -{ - return static_cast(s * c / (S + 1e-20)); -} - -__device__ __forceinline__ float routingWeight(float s, float S, RoutingScale const& scale) -{ - if (scale.fastPath != 0 && S >= 0x1p-13f && s >= 0x1p-60f) - { - float const c = scale.valueFloat; - float const pHi = __fmul_rn(s, c); - float const pLo = __fmaf_rn(s, c, -pHi); - float const q = __fdiv_rn(pHi, S); - float const r = __fmaf_rn(-q, S, pHi); - float const e = __fdiv_rn(__fadd_rn(r, pLo), S); - float const y = __fadd_rn(q, e); - uint32_t const bits = __float_as_uint(y); - if (y > 0.F && (bits & 0x7FFFFFU) != 0 && (bits & 0x7F800000U) != 0x7F800000U) - { - float const ulp = __uint_as_float(bits & 0x7F800000U) * 0x1p-23f; - float const g = __fsub_rn(e, __fsub_rn(y, q)); - if (fabsf(g) < 0.5F * ulp - 0x1p-20f * ulp) - { - return y; - } - } - } - return routingWeightFp64(s, S, scale.value); -} - -__device__ __forceinline__ uint32_t smemAddress(void const* ptr) -{ - return static_cast(__cvta_generic_to_shared(ptr)); -} - -__device__ __forceinline__ void mbarrierInit(uint64_t* bar, uint32_t count) -{ - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(smemAddress(bar)), "r"(count) : "memory"); -} - -// Makes the barrier initialization visible to the bulk-copy (async) proxy. -__device__ __forceinline__ void fenceBarrierInit() -{ - asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); -} - -__device__ __forceinline__ void mbarrierArriveExpectTx(uint64_t* bar, uint32_t bytes) -{ - asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(smemAddress(bar)), "r"(bytes) - : "memory"); -} - -// Arrives on the barrier with release semantics: this thread's earlier shared-memory reads are done before the phase -// can complete. -__device__ __forceinline__ void mbarrierArrive(uint64_t* bar) -{ - asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];\n" ::"r"(smemAddress(bar)) : "memory"); -} - -// Whether the barrier's phase with the given parity has completed (one try, which may suspend the thread for a -// while). -__device__ __forceinline__ bool mbarrierTryWait(uint64_t* bar, uint32_t parity) -{ - uint32_t done = 0; - asm volatile( - "{\n" - ".reg .pred complete;\n" - "mbarrier.try_wait.parity.shared::cta.b64 complete, [%1], %2;\n" - "selp.u32 %0, 1, 0, complete;\n" - "}\n" - : "=r"(done) - : "r"(smemAddress(bar)), "r"(parity) - : "memory"); - return done != 0; -} - -__device__ __forceinline__ void mbarrierWait(uint64_t* bar, uint32_t parity) -{ - while (!mbarrierTryWait(bar, parity)) - { - } -} - -// Waits at named barrier `id` until `threads` threads (whole warps) have arrived. -__device__ __forceinline__ void namedBarrierSync(int id, int threads) -{ - asm volatile("bar.sync %0, %1;\n" ::"r"(id), "r"(threads) : "memory"); -} - -__device__ __forceinline__ uint64_t globalTimerNs() -{ - uint64_t now; - asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(now)); - return now; -} - -// Streamed weights are read once per call: mark their L2 lines evict-first. -__device__ __forceinline__ uint64_t evictFirstPolicy() -{ - uint64_t policy; - asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;\n" : "=l"(policy)); - return policy; -} - -// 1D bulk copy global -> shared that completes `bytes` transactions on `bar`. -__device__ __forceinline__ void bulkCopyToShared( - void* dst, void const* src, uint32_t bytes, uint64_t* bar, uint64_t policy) -{ - uint64_t const srcGlobal = static_cast(__cvta_generic_to_global(src)); - asm volatile( - "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], %4;\n" - : - : "r"(smemAddress(dst)), "l"(srcGlobal), "r"(bytes), "r"(smemAddress(bar)), "l"(policy) - : "memory"); -} - -__device__ __forceinline__ void mmaBf16( - float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) -{ - asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " - "{%0, %1, %2, %3};\n" - : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) - : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); -} - -// One 32-element K step of a 16-row tile times a vector. Lane (g = lane / 4, t = lane % 4) passes elements [8t, 8t + 8) -// of the step from tile row g (lo), tile row g + 8 (hi) and the vector (vec). The fragment's K columns -// {2t, 2t + 1, 2t + 8, 2t + 9} take the lane's elements {0, 1, 2, 3} in the first call and {4, 5, 6, 7} in the second: -// one K permutation for the matrix and the vector, which leaves the dot products unchanged. Every column of the B -// fragment is the vector, so all accumulator elements hold dot products: elements 0 and 1 of row g, 2 and 3 of row -// g + 8. -__device__ __forceinline__ void mmaStep(float (&acc)[4], uint4 const& lo, uint4 const& hi, uint4 const& vec) -{ - mmaBf16(acc, lo.x, hi.x, lo.y, hi.y, vec.x, vec.y); - mmaBf16(acc, lo.z, hi.z, lo.w, hi.w, vec.z, vec.w); -} - -__device__ __forceinline__ float roundToBf16(float value) -{ - return __bfloat162float(__float2bfloat16_rn(value)); -} - -// The FP32 values of the eight BF16 elements of a 16-byte piece, in memory order. -__device__ __forceinline__ void unpackPiece(uint4 const& piece, float (&values)[kPieceElems]) -{ - uint32_t const words[4] = {piece.x, piece.y, piece.z, piece.w}; -#pragma unroll - for (int i = 0; i < 4; ++i) - { - values[2 * i] = __uint_as_float(words[i] << 16); - values[2 * i + 1] = __uint_as_float(words[i] & 0xFFFF0000U); - } -} - -// Dot product of a 16-byte piece of a BF16 weight row with the FP32 values of the same piece of x, in element order -// with FP32 fused multiply-adds. -__device__ __forceinline__ float dotPiece(uint4 const& piece, float const (&values)[kPieceElems]) -{ - float weights[kPieceElems]; - unpackPiece(piece, weights); - float dot = weights[0] * values[0]; -#pragma unroll - for (int e = 1; e < kPieceElems; ++e) - { - dot = fmaf(weights[e], values[e], dot); - } - return dot; -} - -// The warp sums of two values, each by an xor butterfly over offsets 16, 8, 4, 2, 1, the two shuffle chains -// interleaved. -__device__ __forceinline__ void warpSum2(float& a, float& b) -{ -#pragma unroll - for (int offset = 16; offset > 0; offset /= 2) - { - a += __shfl_xor_sync(0xffffffffU, a, offset); - b += __shfl_xor_sync(0xffffffffU, b, offset); - } -} -#endif - -// FC1 + SwiGLU, one CTA per SM (the grid is the SM count). The work is kTopK * interSize units; unit k * interSize + j -// is intermediate index j of routed slot k, that is the expert's gate row (W1, w3w1 row interSize + j) and up row (W3, -// row j). CTA b takes units [b * U / G, (b + 1) * U / G) and streams their rows, gate then up, each as one bulk copy -// into a ring of numStages row buffers. The producer warp selects the experts and issues the rows; it refills a buffer -// as soon as every consumer warp has read its pieces of it, with no CTA-wide barrier per row. The consumer warps wait -// at a named barrier while the producer selects the experts and issues the first rows. Each consumer thread keeps its -// pieces of x in registers and adds its shares of the dot products of two rows at a time, so that the latency chains of -// the two overlap; each consumer warp reduces its shares with shuffles into shared memory, and each unit's SwiGLU adds -// the 8 warp sums in a fixed order. -template -__global__ void __launch_bounds__(kFc1Threads, 2) - lowMBf16MoeFc1Kernel(Bf16* __restrict__ act, Bf16 const* __restrict__ x, float const* __restrict__ routerLogits, - BiasT const* __restrict__ bias, Bf16 const* __restrict__ w3w1, int hiddenSize, int interSize, int numStages) -{ -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - extern __shared__ __align__(128) uint8_t smem[]; - // rowFull[s]: buffer s holds its next row (the producer's arrival plus the copy's bytes). rowFree[s]: every - // consumer warp has read its pieces of buffer s. - __shared__ uint64_t rowFull[kMaxStages]; - __shared__ uint64_t rowFree[kMaxStages]; - __shared__ int32_t topExperts[kTopK]; - - int const warp = threadIdx.x / 32; - int const lane = threadIdx.x % 32; - int const rowBytes = hiddenSize * static_cast(sizeof(Bf16)); - int const numUnits = kTopK * interSize; - int const unitBegin = static_cast(static_cast(blockIdx.x) * numUnits / gridDim.x); - int const unitEnd = static_cast(static_cast(blockIdx.x + 1) * numUnits / gridDim.x); - int const numRows = 2 * (unitEnd - unitBegin); - uint8_t* ring = smem; - float* rowSums = reinterpret_cast(smem + numStages * rowBytes); - - if (threadIdx.x == 0) - { - for (int s = 0; s < numStages; ++s) - { - mbarrierInit(&rowFull[s], 1); - mbarrierInit(&rowFree[s], kFc1ConsumerWarps); - } - fenceBarrierInit(); - } - // The correction bias is a constant: the producer reads it before the grid-dependency wait. - float biasValues[kCandidates] = {}; - if (warp == kFc1ProducerWarp) - { - loadBias(bias, lane, biasValues); - } - __syncthreads(); - - // x and the router logits are written by earlier kernels. - cudaGridDependencySynchronize(); - // The FC2 kernel reads only the logits, the bias and its weights before its own wait, and the logits are complete - // since the wait above returned. Released here, its CTAs take the room beside this kernel's and stage their - // weights while this kernel streams. - cudaTriggerProgrammaticLaunchCompletion(); - - if (warp == kFc1ProducerWarp) - { - TopKKey keys[kCandidates]; - buildKeys(routerLogits, biasValues, lane, keys, nullptr); - int32_t const expert = selectTopK(keys, lane); - if (lane < kTopK) - { - topExperts[lane] = expert; - } - __syncwarp(); - - // Lane 0 streams the rows. The next row is of unit (slot, index), the gate row first, into buffer s, which it - // fills for the pass-th time. - uint64_t const policy = evictFirstPolicy(); - int64_t const expertStride = 2 * static_cast(interSize) * hiddenSize; - int slot = unitBegin / interSize; - int index = unitBegin % interSize; - bool gateRow = true; - Bf16 const* expertRows = w3w1 + topExperts[slot] * expertStride; - int s = 0; - uint32_t pass = 0; - auto issueRow = [&]() - { - int64_t const rowInExpert = gateRow ? interSize + index : index; - mbarrierArriveExpectTx(&rowFull[s], rowBytes); - bulkCopyToShared(ring + s * rowBytes, expertRows + rowInExpert * hiddenSize, rowBytes, &rowFull[s], policy); - if (!gateRow && ++index == interSize) - { - // The up row closed the slot's last unit. - index = 0; - if (++slot < kTopK) - { - expertRows = w3w1 + topExperts[slot] * expertStride; - } - } - gateRow = !gateRow; - if (++s == numStages) - { - s = 0; - ++pass; - } - }; - int const firstRows = numStages < numRows ? numStages : numRows; - if (lane == 0) - { - for (int row = 0; row < firstRows; ++row) - { - issueRow(); - } - } - // Release the consumer warps. - __syncwarp(); - namedBarrierSync(kFc1ParkBarrier, kFc1Threads); - if (lane == 0) - { - for (int row = firstRows; row < numRows; ++row) - { - // Buffer s still holds row - numStages until every consumer warp has read it. - mbarrierWait(&rowFree[s], (pass - 1) & 1U); - issueRow(); - } - } - } - else - { - float xValues[kMaxPiecesPerThread][kPieceElems] = {}; - int const piecesPerThread = hiddenSize / kHiddenGranule; -#pragma unroll - for (int i = 0; i < kMaxPiecesPerThread; ++i) - { - if (i < piecesPerThread) - { - unpackPiece( - __ldcg(reinterpret_cast(x) + threadIdx.x + i * kFc1ConsumerThreads), xValues[i]); - } - } - // Wait at the barrier, without polling the ring, while the producer selects the experts and issues the first - // rows. - namedBarrierSync(kFc1ParkBarrier, kFc1Threads); - int s0 = 0; - uint32_t pass0 = 0; - // Rows row and row + 1 (numRows is even) are in buffers s0 and s1 (numStages >= 2), filled for the pass0-th and - // pass1-th time. - for (int row = 0; row < numRows; row += 2) - { - int s1 = s0 + 1; - uint32_t pass1 = pass0; - if (s1 == numStages) - { - s1 = 0; - ++pass1; - } - mbarrierWait(&rowFull[s0], pass0 & 1U); - mbarrierWait(&rowFull[s1], pass1 & 1U); - uint4 const* pieces0 = reinterpret_cast(ring + s0 * rowBytes); - uint4 const* pieces1 = reinterpret_cast(ring + s1 * rowBytes); - float dot0 = dotPiece(pieces0[threadIdx.x], xValues[0]); - float dot1 = dotPiece(pieces1[threadIdx.x], xValues[0]); -#pragma unroll - for (int i = 1; i < kMaxPiecesPerThread; ++i) - { - if (i < piecesPerThread) - { - dot0 += dotPiece(pieces0[threadIdx.x + i * kFc1ConsumerThreads], xValues[i]); - dot1 += dotPiece(pieces1[threadIdx.x + i * kFc1ConsumerThreads], xValues[i]); - } - } - // This warp is done with both buffers (the arrival releases its lanes' reads, ordered before it by the - // warp barrier); the producer refills a buffer once every consumer warp is. - __syncwarp(); - if (lane == 0) - { - mbarrierArrive(&rowFree[s0]); - mbarrierArrive(&rowFree[s1]); - } - warpSum2(dot0, dot1); - if (lane == 0) - { - rowSums[row * kFc1ConsumerWarps + warp] = dot0; - rowSums[(row + 1) * kFc1ConsumerWarps + warp] = dot1; - } - s0 = s1 + 1; - pass0 = pass1; - if (s0 == numStages) - { - s0 = 0; - ++pass0; - } - } - } - __syncthreads(); - - for (int unit = static_cast(threadIdx.x); unit < numRows / 2; unit += kFc1Threads) - { - float gate = 0.F; - float up = 0.F; - for (int w = 0; w < kFc1ConsumerWarps; ++w) - { - gate += rowSums[2 * unit * kFc1ConsumerWarps + w]; - up += rowSums[(2 * unit + 1) * kFc1ConsumerWarps + w]; - } - // The CUTLASS FC1 GEMM writes its output in BF16 and the activation kernel reads it back. - gate = roundToBf16(gate); - up = roundToBf16(up); - // silu(gate) * up in FP32 as the CUTLASS activation (GLUAdaptor) computes it, then BF16. - float const silu = gate * (1.F / (1.F + expf(-gate))); - act[unitBegin + unit] = __float2bfloat16_rn(silu * up); - } -#endif -} - -// FC2 + the weighted sum over the routed experts. Grid: one CTA per kTiles * 16 output rows; warp k computes routed -// expert k's terms from its weight slice, staged in shared memory before the grid-dependency wait. Warp 0 routes before -// that wait, from the router logits and the bias staged by bulk copies, without FP64 instructions. After the wait every -// warp stages its expert's SwiGLU outputs itself. -template -__global__ void __launch_bounds__(kFc2Threads, 2) - lowMBf16MoeFc2Kernel(Bf16* __restrict__ out, Bf16 const* __restrict__ act, float const* __restrict__ routerLogits, - BiasT const* __restrict__ bias, Bf16 const* __restrict__ w2, int hiddenSize, int interSize, RoutingScale scale) -{ -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - constexpr int kRows = kTiles * kMmaRows; - extern __shared__ __align__(128) uint8_t smem[]; - __shared__ __align__(16) float logitsShared[kNumExperts]; - __shared__ __align__(16) BiasT biasShared[kNumExperts]; - __shared__ uint64_t inputsFull; - __shared__ uint64_t slicesFull; - __shared__ float sigmoidScores[kNumExperts]; - __shared__ int32_t topExperts[kTopK]; - __shared__ float topWeights[kTopK]; - __shared__ float terms[kTopK][kRows]; - - uint64_t const entryNs = globalTimerNs(); - - int const sliceElems = kRows * interSize; - Bf16* slices = reinterpret_cast(smem); - Bf16* actShared = slices + kTopK * sliceElems; - - int const warp = threadIdx.x / 32; - int const lane = threadIdx.x % 32; - int const row0 = blockIdx.x * kRows; - - if (threadIdx.x == 0) - { - mbarrierInit(&inputsFull, 1); - mbarrierInit(&slicesFull, 1); - fenceBarrierInit(); - } - __syncthreads(); - - if (warp == 0) - { - // Before the grid-dependency wait this kernel reads only the router logits, the bias and the FC2 weights. The - // logits come from the router GEMM, which completed before the FC1 kernel returned from its own wait and - // released this grid; the bias and the weights are constants. The logits and the bias come in by bulk copies: - // plain loads issued now queue behind the FC1 kernel's bulk copies on this SM for microseconds (measured). - bool const bulkInputs - = ((reinterpret_cast(routerLogits) | reinterpret_cast(bias)) % 16) == 0; - if (bulkInputs) - { - if (lane == 0) - { - uint64_t const policy = evictFirstPolicy(); - mbarrierArriveExpectTx(&inputsFull, sizeof(logitsShared) + sizeof(biasShared)); - bulkCopyToShared(logitsShared, routerLogits, sizeof(logitsShared), &inputsFull, policy); - bulkCopyToShared(biasShared, bias, sizeof(biasShared), &inputsFull, policy); - } - mbarrierWait(&inputsFull, 0); - } - else - { - // Each lane stages, and below reads back, only the experts it holds. -#pragma unroll - for (int i = 0; i < kCandidates; ++i) - { - logitsShared[i * 32 + lane] = __ldcg(routerLogits + i * 32 + lane); - biasShared[i * 32 + lane] = bias[i * 32 + lane]; - } - } - // The FC1 kernel's producer warps select the same experts on these SMs at about this time, on the FC1 - // kernel's critical path, and routing here concurrently makes their selection about twice as long (measured - // 0.5 -> 1.1 us). So the routing waits until kFc2RouteHoldNs after this CTA's entry; the weight copies are - // still issued long before the FC1 kernel ends. - while (globalTimerNs() < entryNs + kFc2RouteHoldNs) - { - __nanosleep(64); - } - // The FC1 kernel's routing (the keys of buildKeys, selectTopK) on the staged inputs, and the weights from - // routingWeight: the experts, their order and the FP32 weights of noaux_tc_op. - TopKKey keys[kCandidates]; -#pragma unroll - for (int i = 0; i < kCandidates; ++i) - { - int const expert = i * 32 + lane; - float const sigmoid = sigmoidAccurate(logitsShared[expert]); - sigmoidScores[expert] = sigmoid; - keys[i] = TopKKey{sigmoid + static_cast(biasShared[expert]), expert}; - } - int32_t const expert = selectTopK(keys, lane); - __syncwarp(); - auto const tile = cg::tiled_partition<32>(cg::this_thread_block()); - float const scoreNorm = lane < kTopK ? sigmoidScores[expert] : 0.F; - float const sumNorm = cg::reduce(tile, scoreNorm, cg::plus{}); - if (lane < kTopK) - { - topExperts[lane] = expert; - topWeights[lane] = routingWeight(scoreNorm, sumNorm, scale); - } - __syncwarp(); - if (lane == 0) - { - mbarrierArriveExpectTx(&slicesFull, kTopK * sliceElems * static_cast(sizeof(Bf16))); - } - __syncwarp(); - if (lane < kTopK) - { - bulkCopyToShared(slices + lane * sliceElems, - w2 + (static_cast(topExperts[lane]) * hiddenSize + row0) * interSize, - sliceElems * static_cast(sizeof(Bf16)), &slicesFull, evictFirstPolicy()); - } - } - __syncthreads(); - - // The SwiGLU outputs are written by the FC1 kernel: every warp waits for it, then stages its expert's outputs. - cudaGridDependencySynchronize(); - Bf16 const* slice = slices + warp * sliceElems; - Bf16* expertAct = actShared + warp * interSize; - for (int i = lane; i < interSize / 8; i += 32) - { - reinterpret_cast(expertAct)[i] = __ldcg(reinterpret_cast(act + warp * interSize) + i); - } - __syncwarp(); - mbarrierWait(&slicesFull, 0); - - // Lane (g, t) of warp k accumulates the dot products of slice rows tile * 16 + g and tile * 16 + g + 8 with - // act[k, :], over the two 32-element halves of each 64-element K step in two accumulators. Lanes of odd g take the - // two halves in the other order, so that a warp's loads of 8 rows spread over all 32 banks: column n of the B - // fragments then holds the vector in the order of the rows of n's parity, and the accumulators of those columns - // hold the same products in the same order as with one order for all (which of the two holds which half differs). - int const g = lane / 4; - int const t = lane % 4; - int const swap = g & 1; - float acc[kTiles][2][4] = {}; - for (int j0 = 0; j0 < interSize; j0 += kFc2StepPairK) - { -#pragma unroll - for (int half = 0; half < 2; ++half) - { - int const j = j0 + (half ^ swap) * kStepK + 8 * t; - uint4 const av = *reinterpret_cast(expertAct + j); -#pragma unroll - for (int tileIdx = 0; tileIdx < kTiles; ++tileIdx) - { - uint4 const lo = *reinterpret_cast(slice + (tileIdx * kMmaRows + g) * interSize + j); - uint4 const hi = *reinterpret_cast(slice + (tileIdx * kMmaRows + g + 8) * interSize + j); - mmaStep(acc[tileIdx][half], lo, hi, av); - } - } - } - - // The CUTLASS FC2 epilogue rounds each expert's scaled output to BF16 before adding it into the output. Lane (g, 0) - // holds columns 0 and 1 of rows g and g + 8, of which column g & 1 has their dot products. - float const weight = topWeights[warp]; - if (t == 0) - { -#pragma unroll - for (int tileIdx = 0; tileIdx < kTiles; ++tileIdx) - { - float const lo0 = swap ? acc[tileIdx][0][1] : acc[tileIdx][0][0]; - float const lo1 = swap ? acc[tileIdx][1][1] : acc[tileIdx][1][0]; - float const hi0 = swap ? acc[tileIdx][0][3] : acc[tileIdx][0][2]; - float const hi1 = swap ? acc[tileIdx][1][3] : acc[tileIdx][1][2]; - terms[warp][tileIdx * kMmaRows + g] = roundToBf16(weight * (lo0 + lo1)); - terms[warp][tileIdx * kMmaRows + g + 8] = roundToBf16(weight * (hi0 + hi1)); - } - } - __syncthreads(); - if (threadIdx.x < kRows) - { - float sum = 0.F; -#pragma unroll - for (int k = 0; k < kTopK; ++k) - { - sum += terms[k][threadIdx.x]; - } - out[row0 + threadIdx.x] = __float2bfloat16_rn(sum); - } -#endif -} - -// Lets a kernel use smemBytes of dynamic shared memory and asks for the largest shared-memory carveout for both -// kernels, so that an SM running a CTA of one can take a CTA of the other. -template -void setSharedMemoryAttributes(Kernel* kernel, size_t smemBytes) -{ - TLLM_CUDA_CHECK( - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smemBytes))); - TLLM_CUDA_CHECK(cudaFuncSetAttribute( - kernel, cudaFuncAttributePreferredSharedMemoryCarveout, static_cast(cudaSharedmemCarveoutMaxShared))); -} - -// routingWeight's FP32 path applies when the scaling factor is a float with at most 20 significant bits (the low 4 -// mantissa bits zero) between 2^-20 and 2^20. -RoutingScale makeRoutingScale(double routedScalingFactor) -{ - float const valueFloat = static_cast(routedScalingFactor); - uint32_t bits = 0; - std::memcpy(&bits, &valueFloat, sizeof(bits)); - bool const fastPath = static_cast(valueFloat) == routedScalingFactor && (bits & 0xFU) == 0 - && valueFloat >= 0x1p-20f && valueFloat <= 0x1p20f; - return RoutingScale{routedScalingFactor, valueFloat, fastPath ? 1 : 0}; -} - -template -void launchFc2(Bf16* out, Bf16 const* act, float const* routerLogits, BiasT const* bias, Bf16 const* w2, int hiddenSize, - int interSize, double routedScalingFactor, cudaStream_t stream) -{ - int const rows = kTiles * kMmaRows; - size_t const smemBytes = static_cast(kTopK) * (rows + 1) * interSize * sizeof(Bf16); - auto* kernel = &lowMBf16MoeFc2Kernel; - setSharedMemoryAttributes(kernel, smemBytes); - common::launchWithPdlWhenEnabled("lowMBf16MoeFc2", kernel, dim3(hiddenSize / rows), dim3(kFc2Threads), smemBytes, - stream, out, act, routerLogits, bias, w2, hiddenSize, interSize, makeRoutingScale(routedScalingFactor)); -} - -} // namespace - -bool lowMBf16MoeSupportsShape(int64_t hiddenSize, int64_t interSize) -{ - return hiddenSize > 0 && hiddenSize % kHiddenGranule == 0 && hiddenSize <= kMaxPiecesPerThread * kHiddenGranule - && interSize > 0 && interSize % kFc2StepPairK == 0 && interSize <= kMaxInterSize; -} - -template -void invokeLowMBf16Moe(Bf16* out, Bf16* act, Bf16 const* x, float const* routerLogits, BiasT const* bias, - Bf16 const* w3w1, Bf16 const* w2, int hiddenSize, int interSize, double routedScalingFactor, cudaStream_t stream) -{ - TLLM_CHECK_WITH_INFO(lowMBf16MoeSupportsShape(hiddenSize, interSize), - "lowMBf16Moe: unsupported hidden size %d or intermediate size %d", hiddenSize, interSize); - - // FC1: one CTA per SM, each with an even share of the units and a slot per row and consumer warp for the sums. - static int const smCount = common::getMultiProcessorCount(); - int const numUnits = kTopK * interSize; - int const fc1Grid = std::min(smCount, numUnits); - int const rowBytes = hiddenSize * static_cast(sizeof(Bf16)); - int const numStages = std::min(kMaxStages, kFc1RingBytes / rowBytes); - int const maxRows = 2 * ((numUnits + fc1Grid - 1) / fc1Grid); - size_t const fc1SmemBytes - = static_cast(numStages) * rowBytes + static_cast(maxRows) * kFc1ConsumerWarps * sizeof(float); - auto* fc1 = &lowMBf16MoeFc1Kernel; - setSharedMemoryAttributes(fc1, fc1SmemBytes); - common::launchWithPdlWhenEnabled("lowMBf16MoeFc1", fc1, dim3(fc1Grid), dim3(kFc1Threads), fc1SmemBytes, stream, act, - x, routerLogits, bias, w3w1, hiddenSize, interSize, numStages); - - // Two 16-row tiles per FC2 CTA where the 8 weight slices fit in shared memory, else one. - if (interSize <= kMaxInterSize / 2) - { - launchFc2(out, act, routerLogits, bias, w2, hiddenSize, interSize, routedScalingFactor, stream); - } - else - { - launchFc2(out, act, routerLogits, bias, w2, hiddenSize, interSize, routedScalingFactor, stream); - } -} - -#define INSTANTIATE_LOW_M_BF16_MOE(BiasT) \ - template void invokeLowMBf16Moe(Bf16 * out, Bf16 * act, Bf16 const* x, float const* routerLogits, \ - BiasT const* bias, Bf16 const* w3w1, Bf16 const* w2, int hiddenSize, int interSize, \ - double routedScalingFactor, cudaStream_t stream); - -INSTANTIATE_LOW_M_BF16_MOE(float); -INSTANTIATE_LOW_M_BF16_MOE(__nv_bfloat16); - -#undef INSTANTIATE_LOW_M_BF16_MOE - -} // namespace kernels::glm_moe_dsa - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMoeFusedRouter.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMoeFusedRouter.cu deleted file mode 100644 index 1391d9bedbef..000000000000 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMoeFusedRouter.cu +++ /dev/null @@ -1,574 +0,0 @@ -/* - * 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. - */ - -// Router GEMV + DeepSeek-style routing (sigmoid + bias, top-8 of 256 experts, scaled sum normalization) in one kernel -// for 1-8 tokens, with the outputs of routingIndicesSmallTokenKernel. -// -// Logits. Every logit is the value dsv3RouterGemm.cu's router_gemm_kernel computes (bit for bit). There, thread t -// (0..127) of the CTA of expert e accumulates acc = fma(a, w, acc) over k = ki * 1024 + 8 t + j (ki = 0..5, j = 0..7), -// in that order; warp s (the "slice" of threads 32 s .. 32 s + 31) adds its lanes with the xor butterfly 16, 8, 4, 2, 1 -// (lane 0's value), and the four slice sums are added to 0.0f in slice order. Here a GEMV CTA covers one slice of 8 -// experts: it stages the slice's activations of every token in shared memory (24 KB at 8 tokens), runs the same -// per-thread chains (fma.rn.f32.bf16 computes fma(float(a), float(w), acc) without conversion instructions), adds the -// 32 positions with the same pairs as the butterfly and publishes the slice sums. 128 CTAs cover 256 experts x 4 -// slices, and each CTA stages a quarter of the activations instead of all of them. -// -// Handoff. The slice sums go to one of two exchange buffers in the workspace (by the parity of the epoch word) whose -// idle value is -0.0. A sum of -0.0 is published as +0.0, which gives the same logit, since the four sums are added to -// 0.0f. The routing CTA polls its token rows until no -0.0 is left, so no fence or counter is needed. It resets the -// other buffer, consumed by the previous launch, while it waits, and advances the epoch at its end. Launches that share -// a workspace must therefore run in stream order. -// -// Routing. Phase A of routingIndicesSmallTokenKernel (one warp per token: sigmoid + bias, the packed-key top-K, the -// scaled sum normalization; the same operations in the same order). The metadata then follows from the set of -// selected local experts after a single barrier: with no more tokens than tile rows every selected local expert owns -// exactly one CTA tile, so its tile index is its rank among the selected local experts and its rows start at -// rank * tile, which is what the small-token kernel's scans compute. - -#include "tensorrt_llm/common/envUtils.h" -#include "tensorrt_llm/kernels/glmMoeDsa/glmMoeFusedRouter.h" -#include "tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingCustomPolicy.cuh" -#include "tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingKernel.h" - -#include - -// The kernel is file-local to the routing namespace whose policies and helpers it uses. -namespace moe::dev::routing -{ -namespace routingCustom -{ -namespace -{ - -using FusedExpertSelect = TopKExpertSelect; -using FusedKernelParams = KernelParams; - -constexpr int kNumExperts = 256; -constexpr int kTopK = 8; -constexpr int kMaxTokens = 8; -constexpr int kHiddenDim = 6144; -constexpr int kEltsPerVec = 8; // bf16 per 16-byte load, as in the router GEMV -constexpr int kKIterElts = kEltsPerVec * 128; // elements per k-iteration of the router GEMV -constexpr int kKIters = kHiddenDim / kKIterElts; // 6 -constexpr int kSlices = 4; // warps per expert in the router GEMV -constexpr int kExpertsPerCta = 8; -constexpr int kNumGemvCtas = kNumExperts / kExpertsPerCta * kSlices; // 128 -constexpr int kNumThreads = 256; -constexpr int kSliceVecs = kKIters * WarpSize; // 16-byte vectors of one token's slice -constexpr int kActBytes = kMaxTokens * kSliceVecs * 16; // 24 KB -constexpr int kPartBytes = WarpSize * kExpertsPerCta * kMaxTokens * 4; // 8 KB -constexpr int kSmemBytes = kActBytes + kPartBytes; -constexpr int kExchangeFloats = kMaxTokens * kNumExperts * kSlices; -constexpr int kKIdxRowBytes = 8; -constexpr uint32_t kNegZeroBits = 0x80000000u; - -static_assert(kNumExperts == kNumThreads, "the routing CTA runs one thread per expert"); -static_assert( - 2 * kExchangeFloats + 4 == tensorrt_llm::kernels::glm_moe_dsa::kFusedRouterWorkspaceWords, "workspace layout"); - -struct FusedRouterParams -{ - FusedKernelParams routing; - __nv_bfloat16 const* hiddenStates; - __nv_bfloat16 const* routerWeight; - int32_t* workspace; -}; - -__device__ __forceinline__ uint32_t smemAddr(void const* p) -{ - return static_cast(__cvta_generic_to_shared(p)); -} - -__device__ __forceinline__ void gridDependencyWait() -{ - // The memory clobber keeps the compiler from moving memory accesses across the wait. - asm volatile("griddepcontrol.wait;" ::: "memory"); -} - -// Non-coherent 16-byte load of router weights. -__device__ __forceinline__ uint4 loadWeight16(void const* ptr) -{ - uint4 v; - asm volatile("ld.global.nc.v4.u32 {%0, %1, %2, %3}, [%4];" : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(ptr)); - return v; -} - -__device__ __forceinline__ uint32_t loadVolatile(uint32_t const* ptr) -{ - uint32_t v; - asm volatile("ld.volatile.global.u32 %0, [%1];" : "=r"(v) : "l"(ptr) : "memory"); - return v; -} - -__device__ __forceinline__ float4 loadVolatile4(float const* ptr) -{ - float4 v; - asm volatile("ld.volatile.global.v4.f32 {%0, %1, %2, %3}, [%4];" - : "=f"(v.x), "=f"(v.y), "=f"(v.z), "=f"(v.w) - : "l"(ptr) - : "memory"); - return v; -} - -// A slice sum for the routing CTA, which polls it with volatile loads. -__device__ __forceinline__ void storeRelaxed(float* ptr, float v) -{ - asm volatile("st.relaxed.gpu.global.f32 [%0], %1;" ::"l"(ptr), "f"(v) : "memory"); -} - -__device__ __forceinline__ bool isNegZero(float v) -{ - return __float_as_uint(v) == kNegZeroBits; -} - -// acc = fma(a[j], w[j], acc) for j = 0..7, in that order, on 8 packed bf16 pairs. -__device__ __forceinline__ void fmaBf16x8(float& acc, uint4 const& a, uint4 const& w) -{ - uint32_t const aw[4] = {a.x, a.y, a.z, a.w}; - uint32_t const ww[4] = {w.x, w.y, w.z, w.w}; -#pragma unroll - for (int q = 0; q < 4; q++) - { - asm("{\n" - ".reg .b16 al, ah, bl, bh;\n" - "mov.b32 {al, ah}, %1;\n" - "mov.b32 {bl, bh}, %2;\n" - "fma.rn.f32.bf16 %0, al, bl, %0;\n" - "fma.rn.f32.bf16 %0, ah, bh, %0;\n" - "}\n" - : "+f"(acc) - : "r"(aw[q]), "r"(ww[q])); - } -} - -// The slice GEMV of one CTA: slice r of experts 8 g .. 8 g + 7. Thread mapping: warp w covers positions 8 (w & 3) .. -// + 7 of the slice and experts 4 (w >> 2) .. + 3, so the 4 expert groups of a warp read the same 128 bytes of -// activations. Returns through emit(eLocal, m, sliceSum) once per (expert, token). -template -__device__ __forceinline__ void sliceGemv(uint8_t* smem, __nv_bfloat16 const* hiddenStates, uint4 const (&w)[kKIters], - int r, int t, int eLocal, Emit const& emit) -{ - uint4* act = reinterpret_cast(smem); - float* part = reinterpret_cast(smem + kActBytes); - - // Stage the slice: token m, k-iteration ki, positions 32 r .. 32 r + 31. - constexpr int kVecs = kNumTokens * kSliceVecs; -#pragma unroll - for (int j = 0; j < (kVecs + kNumThreads - 1) / kNumThreads; ++j) - { - int const c = threadIdx.x + kNumThreads * j; - if (c < kVecs) - { - int const m = c / kSliceVecs; - int const rem = c - m * kSliceVecs; - __nv_bfloat16 const* src = hiddenStates + m * kHiddenDim + (rem / WarpSize) * kKIterElts - + (WarpSize * r + rem % WarpSize) * kEltsPerVec; - asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ::"r"(smemAddr(&act[c])), "l"(src) : "memory"); - } - } - asm volatile("cp.async.commit_group;\ncp.async.wait_group 0;" ::: "memory"); - __syncthreads(); - - // The per-thread chains of router thread 32 r + t for one expert and every token. - float acc[kMaxTokens] = {}; -#pragma unroll - for (int ki = 0; ki < kKIters; ki++) - { -#pragma unroll - for (int m = 0; m < kNumTokens; m++) - { - fmaBf16x8(acc[m], act[(m * kKIters + ki) * WarpSize + t], w[ki]); - } - } - - // The butterfly of the slice: 4 threads per (expert, token) fold the positions {sub, sub + 4, ..., sub + 28} - // through levels 16, 8 and 4, two shuffles add levels 2 and 1. - float4* row = reinterpret_cast(part + (t * kExpertsPerCta + eLocal) * kMaxTokens); - row[0] = make_float4(acc[0], acc[1], acc[2], acc[3]); - row[1] = make_float4(acc[4], acc[5], acc[6], acc[7]); - __syncthreads(); - int const pair = threadIdx.x / 4; - int const sub = threadIdx.x % 4; - int const eL = pair / kMaxTokens; - int const m = pair % kMaxTokens; - auto v = [&](int tt) { return part[(tt * kExpertsPerCta + eL) * kMaxTokens + m]; }; - float const s1a = v(sub) + v(sub + 16); - float const s1b = v(sub + 4) + v(sub + 20); - float const s1c = v(sub + 8) + v(sub + 24); - float const s1d = v(sub + 12) + v(sub + 28); - float sum = (s1a + s1c) + (s1b + s1d); - sum += __shfl_xor_sync(0xffffffff, sum, 2); - sum += __shfl_xor_sync(0xffffffff, sum, 1); - if (sub == 0 && m < kNumTokens) - { - emit(eL, m, sum); - } -} - -template -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) -__global__ void __launch_bounds__(kNumThreads, 1) routingFusedRouterKernel(FusedRouterParams const __grid_constant__ p) -{ - extern __shared__ __align__(128) uint8_t smemSlice[]; - __shared__ __align__(8) uint8_t smemKIdx[kNumExperts * kKIdxRowBytes]; - __shared__ int16_t smemTopKExpert[kMaxTokens * kTopK]; - __shared__ float smemBias[kNumExperts]; - __shared__ __align__(16) uint32_t smemSelected[kNumExperts / 32]; - - auto const& params = p.routing; - float* const buffers = reinterpret_cast(p.workspace); - uint32_t* const epochPtr = reinterpret_cast(p.workspace + 2 * kExchangeFloats); - - if (blockIdx.x < kNumGemvCtas) - { - int const r = blockIdx.x % kSlices; - int const g = blockIdx.x / kSlices; - int const warp = threadIdx.x / WarpSize; - int const lane = threadIdx.x % WarpSize; - int const t = 8 * (warp % 4) + lane % 8; - int const eLocal = 4 * (warp / 4) + lane / 8; - // The weights are not written by the primary grid: load them before the wait. A shared-memory store of their - // XOR ahead of a barrier makes the loads complete before the wait instead of being scheduled after it. - uint4 w[kKIters]; - __nv_bfloat16 const* wRow - = p.routerWeight + (kExpertsPerCta * g + eLocal) * kHiddenDim + (WarpSize * r + t) * kEltsPerVec; - uint32_t loaded = 0; -#pragma unroll - for (int ki = 0; ki < kKIters; ki++) - { - w[ki] = loadWeight16(wRow + ki * kKIterElts); - loaded ^= w[ki].x ^ w[ki].y ^ w[ki].z ^ w[ki].w; - } - reinterpret_cast(smemSlice + kActBytes)[threadIdx.x] = loaded; - __syncthreads(); - gridDependencyWait(); - // The previous launch on this workspace precedes the primary grid in the stream, so it has completed. - float* const exchange = buffers + (loadVolatile(epochPtr) & 1u) * kExchangeFloats; - sliceGemv(smemSlice, p.hiddenStates, w, r, t, eLocal, - [&](int eL, int m, float sliceSum) - { - storeRelaxed(&exchange[(m * kNumExperts + kExpertsPerCta * g + eL) * kSlices + r], - isNegZero(sliceSum) ? 0.f : sliceSum); - }); - // Exit without a trigger: the dependents launch once the routing CTA triggers. - return; - } - - // Routing CTA. Independent of the primary grid: clear this expert's row of the byte map, stage its bias. - *reinterpret_cast(&smemKIdx[threadIdx.x * kKIdxRowBytes]) = uint2{0xFFFFFFFFu, 0xFFFFFFFFu}; - auto const& pre = params.mExpertSelectParams.mPreprocessParams; - smemBias[threadIdx.x] = loadScalar(pre.ptrRoutingBias, threadIdx.x, pre.dtypeBias); - if (threadIdx.x < kNumExperts / 32) - { - smemSelected[threadIdx.x] = 0u; - } - __syncthreads(); - gridDependencyWait(); - uint32_t const epoch = loadVolatile(epochPtr); - float* const exchange = buffers + (epoch & 1u) * kExchangeFloats; - { - // The buffer of the previous launch, which its routing CTA has consumed: idle again for the next launch. - float4* other = reinterpret_cast(buffers + ((epoch + 1u) & 1u) * kExchangeFloats); - for (int i = threadIdx.x; i < kExchangeFloats / 4; i += kNumThreads) - { - other[i] = make_float4(-0.f, -0.f, -0.f, -0.f); - } - } - - // ----- Phase A: one warp per token, as in routingIndicesSmallTokenKernel ----- - int32_t const warpIdx = __shfl_sync(0xffffffff, threadIdx.x / WarpSize, 0); - int32_t const laneIdx = cutlass::arch::LaneId(); - auto block = cg::this_thread_block(); - auto warp = cg::tiled_partition(block); - if (warpIdx < kNumTokens) - { - constexpr int VecSize = kNumExperts / WarpSize; - float4 slices[VecSize]; - bool ready; - do - { - ready = true; -#pragma unroll - for (int i = 0; i < VecSize; ++i) - { - slices[i] = loadVolatile4(exchange + (warpIdx * kNumExperts + i * WarpSize + laneIdx) * kSlices); - ready = ready && !isNegZero(slices[i].x) && !isNegZero(slices[i].y) && !isNegZero(slices[i].z) - && !isNegZero(slices[i].w); - } - } while (!ready); - - float score[VecSize]; - int32_t idx[VecSize]; -#pragma unroll - for (int i = 0; i < VecSize; ++i) - { - idx[i] = i * WarpSize + laneIdx; - float logit = 0.0f; - logit += slices[i].x; - logit += slices[i].y; - logit += slices[i].z; - logit += slices[i].w; - score[i] = sigmoid_accurate(logit) + smemBias[idx[i]]; - } - float topScore[kTopK]; - int32_t topExpert[kTopK]; - topk::reduceTopK(warp, topScore, topExpert, score, idx, float{-INFINITY}, kTopK); - - // Lane k's top-K entry, the value the small-token kernel reads as warpTopKScore[laneIdx]. - float myScore = 0.f; - int32_t myExpert = 0; -#pragma unroll - for (int k = 0; k < kTopK; ++k) - { - myScore = laneIdx == k ? topScore[k] : myScore; - myExpert = laneIdx == k ? topExpert[k] : myExpert; - } - bool const isTop = laneIdx < kTopK; - float const sigmoidScore = isTop ? (myScore - smemBias[myExpert]) : 0.f; - // cg::reduce's butterfly (xor 16, 8, 4, 2, 1) adds +0.0 to lanes 0-7 in its first two levels, which leaves - // them unchanged (no sigmoidScore is -0.0); levels 4, 2 and 1 then pair lanes 0-7 as here. - float sum = sigmoidScore; - sum += __shfl_xor_sync(0xffffffff, sum, 4); - sum += __shfl_xor_sync(0xffffffff, sum, 2); - sum += __shfl_xor_sync(0xffffffff, sum, 1); - if (isTop) - { - auto const& post = params.mExpertSelectParams.mPostprocessParams; - float const weight = sigmoidScore * post.routeScale / (sum + post.sumEpsilon); - smemKIdx[myExpert * kKIdxRowBytes + warpIdx] = static_cast(laneIdx); - smemTopKExpert[warpIdx * kTopK + laneIdx] = static_cast(myExpert); - atomicOr(&smemSelected[myExpert / 32], 1u << (myExpert % 32)); - params.mPtrTopKWeights[warpIdx * kTopK + laneIdx] = __nv_bfloat16{weight}; - } - } - __syncthreads(); - - // ----- Metadata: tile index = rank among the selected local experts ----- - int32_t const localBegin = params.mLocalExpertsStartIdx; - int32_t const localEnd = params.mLocalExpertsStartIdx + params.mNumLocalExperts; - uint4 const lo = *reinterpret_cast(&smemSelected[0]); - uint4 const hi = *reinterpret_cast(&smemSelected[4]); - uint32_t selected[8] = {lo.x, lo.y, lo.z, lo.w, hi.x, hi.y, hi.z, hi.w}; - if (localBegin != 0 || localEnd < kNumExperts) - { -#pragma unroll - for (int i = 0; i < 8; ++i) - { - int32_t const b0 = max(localBegin - 32 * i, 0); - int32_t const b1 = min(localEnd - 32 * i, 32); - uint32_t const upTo = b1 >= 32 ? 0xFFFFFFFFu : ((1u << max(b1, 0)) - 1u); - uint32_t const below = b0 >= 32 ? 0xFFFFFFFFu : ((1u << b0) - 1u); - selected[i] &= b1 <= b0 ? 0u : (upTo & ~below); - } - } - int32_t numSelected = 0; -#pragma unroll - for (int i = 0; i < 8; ++i) - { - numSelected += __popc(selected[i]); - } - // (rank, is selected and local) of an expert, from masked popcounts of every word: indexing the words by the expert - // would place them in local memory. - auto rankOf = [&](int32_t e, bool& isSelected) - { - int32_t const word = e >> 5; - uint32_t const bit = 1u << (e & 31); - int32_t rank = 0; - int32_t hit = 0; -#pragma unroll - for (int i = 0; i < 8; ++i) - { - rank += __popc(selected[i] & (i < word ? 0xFFFFFFFFu : (i == word ? bit - 1u : 0u))); - hit += __popc(selected[i] & (i == word ? bit : 0u)); - } - isSelected = hit != 0; - return rank; - }; - auto tileRows = [&](int32_t n) - { return params.mIsPow2 ? mulLog2(n, params.mPaddingLog2) : mulTileN(n, params.mTileTokensDim); }; - - { - int32_t const expert = threadIdx.x; - bool isSelected; - int32_t const rank = rankOf(expert, isSelected); - if (isSelected) - { - uint2 const kRow = *reinterpret_cast(&smemKIdx[expert * kKIdxRowBytes]); - int32_t const count = __popc(__vsetne4(kRow.x, 0xFFFFFFFFu)) + __popc(__vsetne4(kRow.y, 0xFFFFFFFFu)); - params.mPtrCtaIdxXyToBatchIdx[rank] = expert - localBegin; - params.mPtrCtaIdxXyToMnLimit[rank] = min(tileRows(rank + 1), tileRows(rank) + count); - } - } - if (threadIdx.x == 0) - { - params.mPtrPermutedIdxSize[0] = tileRows(numSelected); - params.mPtrNumNonExitingCtas[0] = numSelected; - } - - // Permutation, one (token, k) pair per thread. offset-within-expert = number of earlier tokens that also selected - // the expert (its byte-map row below byte t), the small-token kernel's order. - if (static_cast(threadIdx.x) < kNumTokens * kTopK) - { - int32_t const expandedIdx = threadIdx.x; - int32_t const tokenIdx = expandedIdx / kTopK; - int32_t const pairExpert = smemTopKExpert[expandedIdx]; - bool pairIsLocal; - int32_t const rank = rankOf(pairExpert, pairIsLocal); - uint2 const row = *reinterpret_cast(&smemKIdx[pairExpert * kKIdxRowBytes]); - uint64_t const present = (static_cast(__vsetne4(row.y, 0xFFFFFFFFu)) << 32) - | static_cast(__vsetne4(row.x, 0xFFFFFFFFu)); - uint64_t const earlier = tokenIdx == 0 ? 0ull : (present & ((1ull << (8 * tokenIdx)) - 1)); - int32_t const permutedIdx = pairIsLocal ? tileRows(rank) + __popcll(earlier) : int32_t{-1}; - if (params.mPtrExpandedIdxToPermutedIdx != nullptr) - { - params.mPtrExpandedIdxToPermutedIdx[expandedIdx] = permutedIdx; - } - if (params.mPtrPermutedIdxToExpandedIdx != nullptr && pairIsLocal) - { - params.mPtrPermutedIdxToExpandedIdx[permutedIdx] = expandedIdx; - } - if (params.mPtrPermutedIdxToTokenIdx != nullptr && pairIsLocal) - { - params.mPtrPermutedIdxToTokenIdx[permutedIdx] = tokenIdx; - } - } - if (threadIdx.x == 0) - { - *epochPtr = epoch + 1u; - } - // After all routing outputs are written. - cudaTriggerProgrammaticLaunchCompletion(); -} -#else -__global__ void __launch_bounds__(kNumThreads, 1) - routingFusedRouterKernel(FusedRouterParams const __grid_constant__ /* p */) -{ - assert(false && "routingFusedRouterKernel is only supported on SM100+ architectures"); -} -#endif // if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) - -template -void launchFusedRouter(FusedRouterParams const& params, bool usePdl, cudaStream_t stream) -{ - cudaLaunchConfig_t config{}; - config.gridDim = kNumGemvCtas + 1; - config.blockDim = kNumThreads; - config.dynamicSmemBytes = kSmemBytes; - config.stream = stream; - cudaLaunchAttribute attributes[1] = {}; - attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attributes[0].val.programmaticStreamSerializationAllowed = int(usePdl); - config.attrs = attributes; - config.numAttrs = 1; - TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, routingFusedRouterKernel, params)); -} - -bool supportsFusedRouter(Data const& data, int32_t hiddenDim) -{ - return hiddenDim == kHiddenDim && data.mNumExperts == kNumExperts && data.mTopK == kTopK && data.mNumTokens >= 1 - && data.mNumTokens <= kMaxTokens && data.mNumTokens <= data.mTileTokensDim && data.mLocalExpertsStrideLog2 == 0 - && data.mPreprocessType == RoutingPreprocessType::SigmoidBias - && data.mPostprocessType == RoutingPostprocessType::ScaledSumNormalize - && data.mDtypeOutput == tg::Dtype::Bfloat16 && data.mPtrTopKWeights != nullptr - && data.mPtrRoutingBias != nullptr && data.mPtrTopKIds == nullptr; -} - -void runFusedRouter(Data const& data, void const* hiddenStates, void const* routerWeight, int32_t hiddenDim, - int32_t* workspace, void* stream) -{ - TLLM_CHECK_WITH_INFO(supportsFusedRouter(data, hiddenDim), - "The fused router + routing kernel supports 1-8 tokens (at most the tile size), 256 experts, top-8, hidden " - "size %d and SigmoidBias + ScaledSumNormalize routing; got %d tokens, tile %d, %d experts, top-%d, hidden %d", - kHiddenDim, data.mNumTokens, data.mTileTokensDim, data.mNumExperts, data.mTopK, hiddenDim); - FusedRouterParams params; - params.routing = FusedKernelParams::setKernelParams(data); - params.hiddenStates = static_cast<__nv_bfloat16 const*>(hiddenStates); - params.routerWeight = static_cast<__nv_bfloat16 const*>(routerWeight); - params.workspace = workspace; - auto const s = static_cast(stream); - switch (data.mNumTokens) - { - case 1: launchFusedRouter<1>(params, data.mUsePdl, s); break; - case 2: launchFusedRouter<2>(params, data.mUsePdl, s); break; - case 3: launchFusedRouter<3>(params, data.mUsePdl, s); break; - case 4: launchFusedRouter<4>(params, data.mUsePdl, s); break; - case 5: launchFusedRouter<5>(params, data.mUsePdl, s); break; - case 6: launchFusedRouter<6>(params, data.mUsePdl, s); break; - case 7: launchFusedRouter<7>(params, data.mUsePdl, s); break; - default: launchFusedRouter<8>(params, data.mUsePdl, s); break; - } -} - -} // namespace -} // namespace routingCustom -} // namespace moe::dev::routing - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels::glm_moe_dsa -{ - -void invokeFusedRouter(void const* routerHiddenStates, void const* routerWeight, int32_t hiddenDim, - int32_t* routerWorkspace, void* routingBias, batchedGemm::trtllm::gen::Dtype dtypeRoutingBias, int32_t numTokens, - int32_t numExperts, int32_t topK, int32_t localExpertOffset, int32_t localNumExperts, float routedScalingFactor, - int32_t tileTokensDim, int32_t* permutedIdxSize, int32_t* expandedIdxToPermutedIdx, int32_t* permutedIdxToTokenIdx, - void* expertWeights, int32_t* ctaIdxXyToBatchIdx, int32_t* ctaIdxXyToMnLimit, int32_t* numNonExitingCtas, - cudaStream_t stream) -{ - int32_t paddingLog2 = 0; - while ((1 << paddingLog2) < tileTokensDim) - { - ++paddingLog2; - } - TLLM_CHECK_WITH_INFO( - (1 << paddingLog2) == tileTokensDim, "tileTokensDim must be a power of 2, got %d", tileTokensDim); - - // The routing configuration of the TRTLLM-Gen MoE runner's DeepSeek-V3 routing without groups. - ::moe::dev::routing::routingCustom::Data routingData; - routingData.mDtypeOutput = batchedGemm::trtllm::gen::Dtype::Bfloat16; - routingData.mDtypeInput = batchedGemm::trtllm::gen::Dtype::Fp32; - routingData.mUsePdl = tensorrt_llm::common::getEnvEnablePDL(); - routingData.mPreprocessType = ::moe::dev::routing::RoutingPreprocessType::SigmoidBias; - routingData.mPostprocessType = ::moe::dev::routing::RoutingPostprocessType::ScaledSumNormalize; - routingData.mPtrRoutingBias = routingBias; - routingData.mDtypeBias = dtypeRoutingBias; - routingData.mRouteScale = routedScalingFactor; - routingData.mPtrPermutedIdxSize = permutedIdxSize; - routingData.mPtrExpandedIdxToPermutedIdx = expandedIdxToPermutedIdx; - routingData.mPtrPermutedIdxToExpandedIdx = nullptr; - routingData.mPtrPermutedIdxToTokenIdx = permutedIdxToTokenIdx; - routingData.mPtrTopKWeights = expertWeights; - routingData.mPtrCtaIdxXyToBatchIdx = ctaIdxXyToBatchIdx; - routingData.mPtrCtaIdxXyToMnLimit = ctaIdxXyToMnLimit; - routingData.mPtrNumNonExitingCtas = numNonExitingCtas; - routingData.mNumTokens = numTokens; - routingData.mNumExperts = numExperts; - routingData.mTopK = topK; - routingData.mPaddingLog2 = paddingLog2; - routingData.mTileTokensDim = tileTokensDim; - routingData.mLocalExpertsStartIdx = localExpertOffset; - routingData.mLocalExpertsStrideLog2 = 0; - routingData.mNumLocalExperts = localNumExperts; - - ::moe::dev::routing::routingCustom::runFusedRouter( - routingData, routerHiddenStates, routerWeight, hiddenDim, routerWorkspace, stream); -} - -} // namespace kernels::glm_moe_dsa - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.cu deleted file mode 100644 index 1b5a83166840..000000000000 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.cu +++ /dev/null @@ -1,244 +0,0 @@ -/* - * 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/glmMoeDsa/glmRouterGemm.h" - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/envUtils.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels::glm_moe_dsa -{ - -namespace -{ - -// Convert 8 bfloat16 values from a uint4 to float array - optimized conversion -template -__device__ __forceinline__ void bf16_uint4_to_float8(uint4 const& vec, float* dst) -{ - __nv_bfloat16* bf16_ptr = reinterpret_cast<__nv_bfloat16*>(const_cast(&vec)); - -#pragma unroll - for (int i = 0; i < VPT; i++) - { - dst[i] = __bfloat162float(bf16_ptr[i]); - } -} - -// mat_a: the activations [kNumTokens, kHiddenDim]; mat_b: the weight [kNumExperts, kHiddenDim]; out: the logits -// [kNumTokens, kNumExperts]. Before the grid-dependency wait each CTA prefetches its weight row into L2 (no kernel of -// the step writes the weight), so the weight loads after the wait hit L2; the dependent grid may launch right after the -// wait instead of after the logits are written, and must wait on its own grid dependency before reading them. -template -__global__ __launch_bounds__(128, 1) void router_gemm_kernel(float* out, T const* mat_a, T const* mat_b) -{ - // Each block handles one expert column - int const n_idx = blockIdx.x; - int const tid = threadIdx.x; - constexpr int kWarpSize = 32; - constexpr int kNumWarps = kBlockSize / kWarpSize; - // Constants for this kernel - constexpr int k_elems_per_k_iteration = VPT * kBlockSize; - constexpr int k_iterations = kHiddenDim / k_elems_per_k_iteration; // Total K iterations - - // Initialize accumulators for all M rows - float acc[kNumTokens] = {}; - - // Shared memory for warp-level reduction - __shared__ float sm_reduction[kNumTokens][kNumWarps]; // kNumWarps - - // The weight row of expert n_idx (a column of B in column-major order) - T const* b_col = mat_b + n_idx * kHiddenDim; - - // Pre-compute k_base values for each iteration to help compiler optimize - int k_bases[k_iterations]; -#pragma unroll - for (int ki = 0; ki < k_iterations; ki++) - { - k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT; - } - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - { - // The primary grid does not write the weights, so this CTA's column can stream into L2 while the primary - // runs; the weight loads below then hit L2. What is loaded and the accumulation order are unchanged. - constexpr int kLineBytes = 128; - constexpr int kColumnBytes = kHiddenDim * static_cast(sizeof(T)); - static_assert(kColumnBytes % kLineBytes == 0, "the weight column must span whole 128-byte lines"); - char const* b_col_bytes = reinterpret_cast(b_col); - for (int offset = tid * kLineBytes; offset < kColumnBytes; offset += kBlockSize * kLineBytes) - { - asm volatile("prefetch.global.L2 [%0];" ::"l"(b_col_bytes + offset) : "memory"); - } - } - cudaGridDependencySynchronize(); - // Every dependent waits on its own grid dependency before it reads the logits, so it may launch now and overlap - // its prologue with this GEMV. - cudaTriggerProgrammaticLaunchCompletion(); -#endif - - // Process the GEMM in chunks - for (int ki = 0; ki < k_iterations; ki++) - { - int const k_base = k_bases[ki]; - - // Load B matrix values using vector load (8 bf16 values) - uint4 b_vec = *reinterpret_cast(b_col + k_base); - - // Convert B values to float - float b_float[VPT]; - bf16_uint4_to_float8(b_vec, b_float); - -// Process each token -#pragma unroll - for (int m_idx = 0; m_idx < kNumTokens; m_idx++) - { - // Load both rows of A matrix using vector loads - uint4 a_vec = *reinterpret_cast(mat_a + (m_idx * kHiddenDim) + k_base); - - // Convert A values to float - float a_float[VPT]; - bf16_uint4_to_float8(a_vec, a_float); - -// Process elements in this chunk -#pragma unroll - for (int k = 0; k < VPT; k++) - { - float a = a_float[k]; - float b = b_float[k]; - acc[m_idx] += a * b; - } - } - } - - // Perform warp-level reduction - int const warpSize = 32; - int const warpId = tid / warpSize; - int const laneId = tid % warpSize; - - // Register for warp-level reduction results - float warp_result[kNumTokens]; - -#pragma unroll - for (int m_idx = 0; m_idx < kNumTokens; m_idx++) - { - warp_result[m_idx] = acc[m_idx]; - } - -// Perform warp-level reduction using optimized butterfly pattern -#pragma unroll - for (int m = 0; m < kNumTokens; m++) - { - float sum = warp_result[m]; - - // Butterfly reduction pattern - sum += __shfl_xor_sync(0xffffffff, sum, 16); - sum += __shfl_xor_sync(0xffffffff, sum, 8); - sum += __shfl_xor_sync(0xffffffff, sum, 4); - sum += __shfl_xor_sync(0xffffffff, sum, 2); - sum += __shfl_xor_sync(0xffffffff, sum, 1); - - // Only the first thread in each warp stores to shared memory - if (laneId == 0) - { - sm_reduction[m][warpId] = sum; - } - } - - __syncthreads(); - - // Final reduction across warps (only first thread) - if (tid == 0) - { -#pragma unroll - for (int m = 0; m < kNumTokens; m++) - { - float final_sum = 0.0f; - -// Sum across the kNumWarps -#pragma unroll - for (int w = 0; w < kNumWarps; w++) - { - final_sum += sm_reduction[m][w]; - } - - // Write final result - out[m * kNumExperts + n_idx] = final_sum; - } - } -} - -template -void launchRouterGemm( - float* logits, __nv_bfloat16 const* hiddenStates, __nv_bfloat16 const* weight, cudaStream_t stream) -{ - using T = __nv_bfloat16; - constexpr int VPT = 16 / sizeof(T); - constexpr int kBlockSize = 128; - cudaLaunchConfig_t config; - config.gridDim = kRouterNumExperts; - config.blockDim = kBlockSize; - config.dynamicSmemBytes = 0; - config.stream = stream; - cudaLaunchAttribute attrs[1]; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); - config.numAttrs = 1; - config.attrs = attrs; - TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, - router_gemm_kernel, logits, hiddenStates, - weight)); -} - -// Launches the instance of numTokens, one of kNumTokens..kRouterMaxTokens. -template -void dispatchRouterGemm( - float* logits, __nv_bfloat16 const* hiddenStates, __nv_bfloat16 const* weight, int numTokens, cudaStream_t stream) -{ - if constexpr (kNumTokens <= kRouterMaxTokens) - { - if (numTokens == kNumTokens) - { - launchRouterGemm(logits, hiddenStates, weight, stream); - } - else - { - dispatchRouterGemm(logits, hiddenStates, weight, numTokens, stream); - } - } - else - { - TLLM_THROW("The GLM router GEMV takes 1 to %d tokens, got %d", kRouterMaxTokens, numTokens); - } -} - -} // namespace - -void invokeRouterGemm( - float* logits, __nv_bfloat16 const* hiddenStates, __nv_bfloat16 const* weight, int numTokens, cudaStream_t stream) -{ - TLLM_CHECK_WITH_INFO( - numTokens >= 1, "The GLM router GEMV takes 1 to %d tokens, got %d", kRouterMaxTokens, numTokens); - dispatchRouterGemm<1>(logits, hiddenStates, weight, numTokens, stream); -} - -} // namespace kernels::glm_moe_dsa - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.cu b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmDenseMlpSmallM.cu similarity index 99% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.cu rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmDenseMlpSmallM.cu index 2a5bb29cdad2..16719923a0ec 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.cu +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmDenseMlpSmallM.cu @@ -35,7 +35,7 @@ // Programmatic dependent launch: activations are read only after griddepcontrol.wait; the dependents are released at // entry or after the weight copies are issued (Trigger). Each grid fits in one wave of resident CTAs. -#include "tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmDenseMlpSmallM.h" #include #include diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.h b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmDenseMlpSmallM.h similarity index 100% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.h rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmDenseMlpSmallM.h diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.cu b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmDraftSampleGlue.cu similarity index 71% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.cu rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmDraftSampleGlue.cu index 5a936de3083c..6615e27de09c 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.cu +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmDraftSampleGlue.cu @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmDraftSampleGlue.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" @@ -569,6 +569,7 @@ __device__ __forceinline__ float warpSum(float v) return v; } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) __device__ __forceinline__ void clusterArriveRelaxed() { asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); @@ -584,6 +585,8 @@ __device__ __forceinline__ void clusterWait() asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); } +#endif + template __global__ void __launch_bounds__(kTableThreads) draftEhprojTableKernel(DraftEhprojTableParams const p) { @@ -929,6 +932,347 @@ __global__ void __launch_bounds__(kTableThreads) draftEhprojTableKernel(DraftEhp #endif } +// The draft step of draftSampleGlueKernel on a cluster of kWideCtas CTAs (draftSampleGlueWideKernel). Right after +// the grid dependency wait, in every CTA: +// - warps 0..kWideScanWarps-1 scan the CTA's chunk (the c-th of kWideCtas consecutive chunks) of every row of the +// logits shard and push the CTA's best (index, value) of each row into CTA 0's shared memory; CTA 0 combines them +// in cluster-rank order and publishes the rank's pairs; every CTA that needs the tokens polls the pairs itself; +// norm group 0 (warps 0 and 1) of CTA c then normalizes the embedding row of row c's token; CTA 0 writes the tokens +// and the next-new-tokens rows and rotates the flags; +// - the remaining warps work beside the scan and the exchange: CTA 0's kWideClearWarp empties the previous call's +// buffer, its kWideStateWarp writes the next step's state, and norm group kWideHnormGroup of CTA c normalizes +// hidden row c (from a copy that cp.async brings into shared memory). +// Same outputs as draftSampleGlueKernel. +constexpr int kWideCtas = 8; +constexpr int kWideThreads = 512; +constexpr int kWideWarps = kWideThreads / 32; +constexpr int kWideScanWarps = 12; +constexpr int kWideScanThreads = kWideScanWarps * 32; +// Logit vectors of each row a scanning thread has in flight per round, and the rows of a scan pass. +constexpr int kWideScanVecs = 2; +constexpr int kWideScanRows = 4; +constexpr int kWideClearWarp = kWideScanWarps; +constexpr int kWideStateWarp = kWideScanWarps + 1; +// The last two warps. +constexpr int kWideHnormGroup = kNumNormGroups - 1; +// Named barriers: norm group g uses 1 + g (normRowWindow); group 0's warps also meet on its id for the tokens; the +// scanning warps meet on their own. +constexpr int kWideGroup0Barrier = 1; +constexpr int kWideScanBarrier = 1 + kNumNormGroups; +static_assert(kWideThreads == kBlockSize, "the norm groups and slots are those of draftSampleGlueKernel"); +static_assert(kDraftSampleGlueMaxRows <= kWideCtas, "CTA c normalizes row c"); +static_assert(kDraftSampleGlueMaxRows % kWideScanRows == 0, "whole scan passes"); +static_assert(2 * kWideHnormGroup + 2 == kWideWarps && kWideStateWarp < 2 * kWideHnormGroup && kWideScanWarps >= 2, + "the warp roles must not overlap"); +static_assert(kDraftSampleGlueMaxRows <= kWideScanThreads, "one thread per row combines the warps' bests"); + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +// 16 bytes from global into shared memory without a register (cp.async, cached in L2 only). +__device__ __forceinline__ void cpAsync16(void* smem, void const* gmem) +{ + uint32_t const dst = static_cast(__cvta_generic_to_shared(smem)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ::"r"(dst), "l"(gmem) : "memory"); +} + +// Every cp.async of this thread has landed. +__device__ __forceinline__ void cpAsyncWaitAll() +{ + asm volatile("cp.async.wait_all;" ::: "memory"); +} + +#endif + +__global__ void __launch_bounds__(kWideThreads) draftSampleGlueWideKernel(DraftSampleGlueParams const p) +{ + __shared__ __align__(16) __nv_bfloat16 sWeights[2][kDraftSampleGlueHidden]; + // Hidden row c: thread t of the hnorm group copies, and later reads, the vectors normRowWindow assigns to it. + __shared__ uint4 sHidden[kNormVecs * kNormThreads]; + __shared__ Candidate sWarpBest[kDraftSampleGlueMaxRows][kWideScanWarps]; + // CTA 0 only: the CTAs' bests, written by CTA c into row c before the cluster barrier, read after it. + __shared__ Candidate sClusterBest[kWideCtas][kDraftSampleGlueMaxRows]; + __shared__ Candidate sSlots[kMaxSlots]; + __shared__ int sTokens[kDraftSampleGlueMaxRows]; + __shared__ float sGroupSums[kNumNormGroups][2]; + __shared__ uint4 sFlags; + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + int const cta = static_cast(cluster.block_rank()); + int const tid = static_cast(threadIdx.x); + int const warp = tid / 32; + int const lane = tid % 32; + bool const glue = p.glueOut != nullptr; + bool const embedHalf = glue && p.colLo < kDraftSampleGlueHidden; + bool const hiddenHalf = glue && p.colHi > kDraftSampleGlueHidden; + + // The norm weights are constants: copy the window columns before the grid dependency wait. + if (glue) + { + for (int h = 0; h < 2; ++h) + { + if ((h == 0 && !embedHalf) || (h == 1 && !hiddenHalf)) + { + continue; + } + int const lo = max(p.colLo - h * kDraftSampleGlueHidden, 0); + int const hi = min(p.colHi - h * kDraftSampleGlueHidden, kDraftSampleGlueHidden); + auto const* src = reinterpret_cast(h == 0 ? p.enormWeight : p.hnormWeight); + auto* dst = reinterpret_cast(sWeights[h]); + for (int v = lo / kNormVecElts + tid; v < hi / kNormVecElts; v += kWideThreads) + { + dst[v] = src[v]; + } + } + } + __syncthreads(); + + // The logits come from the preceding kernel (the draft LM head). Dependents read the outputs only after their + // own wait. + cudaGridDependencySynchronize(); + cudaTriggerProgrammaticLaunchCompletion(); + + bool waited = false; + if (warp < kWideScanWarps) + { + if (tid == 0) + { + sFlags = reinterpret_cast(p.exchangeFlags)[0]; + } + + // This CTA's chunk of every row, kWideScanRows rows per pass: the loads of a pass's round in flight together. + int const numVecs = p.vocabShard / kNormVecElts; + int const chunk = (numVecs + kWideCtas - 1) / kWideCtas; + int const vecBegin = min(cta * chunk, numVecs); + int const vecEnd = min(vecBegin + chunk, numVecs); + Candidate best[kDraftSampleGlueMaxRows]; +#pragma unroll + for (int r = 0; r < kDraftSampleGlueMaxRows; ++r) + { + best[r] = Candidate{-INFINITY, INT_MAX}; + } +#pragma unroll + for (int r0 = 0; r0 < kDraftSampleGlueMaxRows; r0 += kWideScanRows) + { + if (r0 >= p.numRows) + { + continue; + } + for (int base = vecBegin + tid; base < vecEnd; base += kWideScanVecs * kWideScanThreads) + { + uint4 vecs[kWideScanRows][kWideScanVecs]; +#pragma unroll + for (int r = 0; r < kWideScanRows; ++r) + { +#pragma unroll + for (int k = 0; k < kWideScanVecs; ++k) + { + int const v = base + k * kWideScanThreads; + if (r0 + r < p.numRows && v < vecEnd) + { + vecs[r][k] = reinterpret_cast(p.logits + (r0 + r) * p.logitsStride)[v]; + } + } + } +#pragma unroll + for (int r = 0; r < kWideScanRows; ++r) + { +#pragma unroll + for (int k = 0; k < kWideScanVecs; ++k) + { + int const v = base + k * kWideScanThreads; + if (r0 + r < p.numRows && v < vecEnd) + { + offerVector(best[r0 + r], vecs[r][k], v * kNormVecElts); + } + } + } + } + } +#pragma unroll + for (int r = 0; r < kDraftSampleGlueMaxRows; ++r) + { + if (r < p.numRows) + { + Candidate const c = warpArgmax(best[r]); + if (lane == 0) + { + sWarpBest[r][warp] = c; + } + } + } + namedBarrier(kWideScanBarrier, kWideScanThreads); + if (tid < p.numRows) + { + Candidate ctaBest = sWarpBest[tid][0]; + for (int w = 1; w < kWideScanWarps; ++w) + { + offer(ctaBest, sWarpBest[tid][w].value, sWarpBest[tid][w].index); + } + *cluster.map_shared_rank(&sClusterBest[cta][tid], 0) = ctaBest; + } + // The pushes above become visible to CTA 0's waiter. + clusterArrive(); + + uint4 const flags = sFlags; + uint32_t const curIndex = flags.x; + uint32_t const bytesPerBuffer = flags.z; + int const numSlots = p.numRows * p.tpSize; + // CTA 0 writes the tokens and the next-new-tokens rows; CTA c normalizes the embedding row c. + bool const needTokens = cta == 0 || (embedHalf && cta < p.numRows); + + if (warp == 0 && needTokens) + { + if (cta == 0) + { + clusterWait(); + waited = true; + // The rank's pair of every row (the CTAs' bests in cluster-rank order, the order of their chunks), + // published through the multicast address: slot [row][rank] of the current buffer. + if (lane < p.numRows) + { + Candidate rankBest = sClusterBest[0][lane]; + for (int c = 1; c < kWideCtas; ++c) + { + offer(rankBest, sClusterBest[c][lane].value, sClusterBest[c][lane].index); + } + int const globalIndex = rankBest.index + p.rank * p.vocabShard; + auto* curMcast + = reinterpret_cast(static_cast(p.exchangeMcast) + curIndex * bytesPerBuffer); + curMcast[lane * p.tpSize + p.rank] + = make_uint4(static_cast(globalIndex), __float_as_uint(rankBest.value), 0U, 0U); + } + } + // Poll the local copy until every rank's pair has landed. + auto const* curLocal = static_cast(p.exchangeLocal) + curIndex * bytesPerBuffer; + for (int s = lane; s < numSlots; s += 32) + { + VolatilePackedLoad packet; + do + { + packet = loadPackedVolatile(curLocal + s * kDraftSampleGluePacketBytes); + } while (isEmptyPacket(packet)); + sSlots[s] = Candidate{__uint_as_float(packet.words[1]), static_cast(packet.words[0])}; + } + __syncwarp(); + // Global argmax per row: the ranks' pairs in rank order (the lower rank holds the lower indices). + if (lane < p.numRows) + { + Candidate rowBest = sSlots[lane * p.tpSize]; + for (int r = 1; r < p.tpSize; ++r) + { + offer(rowBest, sSlots[lane * p.tpSize + r].value, sSlots[lane * p.tpSize + r].index); + } + sTokens[lane] = rowBest.index; + if (cta == 0) + { + p.draftTokens[lane * p.numDraftTokens + p.draftStep] = rowBest.index; + } + } + if (cta == 0 && lane == 0) + { + // Every CTA read the flags before it arrived at the cluster barrier: rotate them as + // LamportFlags::update does, the current buffer becoming the dirty one. + auto* flagPtr = reinterpret_cast(p.exchangeFlags); + flagPtr[0] = make_uint4((curIndex + 1) % 3, curIndex, bytesPerBuffer, 1U); + flagPtr[1] = make_uint4(static_cast(numSlots * kDraftSampleGluePacketBytes), 0U, 0U, 0U); + p.exchangeFlags[8] = 0U; + } + } + + if (warp < 2 && needTokens) + { + // Group 0: the tokens of warp 0, then the embedding row of this CTA. + namedBarrier(kWideGroup0Barrier, kNormThreads); + if (cta == 0 && warp == 1 && p.nextNewTokens != nullptr && lane < p.numRows) + { + // Output rows after the last draft step: the last accepted token, then the draft tokens of every + // step. + int32_t* out = p.nextNewTokens + lane * (p.numDraftTokens + 1); + out[0] = p.acceptedTokens[p.batchIndices[lane] * p.acceptedTokensStride + p.numAccepted[lane] - 1]; + for (int j = 0; j < p.numDraftTokens; ++j) + { + out[1 + j] = j == p.draftStep ? sTokens[lane] : p.draftTokens[lane * p.numDraftTokens + j]; + } + } + if (embedHalf && cta < p.numRows) + { + // enorm of the token embedding: an id outside [0, vocabSize) gives a zero row, as the Embedding + // module does. + int64_t const token = sTokens[cta]; + bool const valid = token >= 0 && token < p.vocabSize; + normRowWindow( + valid ? reinterpret_cast(p.embedTable + token * kDraftSampleGlueHidden) : nullptr, + sWeights[0], p.eps, p.colLo, min(p.colHi, kDraftSampleGlueHidden), + p.glueOut + cta * p.glueOutStride, 0, sGroupSums); + } + } + } + else if (warp == kWideClearWarp) + { + if (cta == 0) + { + // Empty the buffer of the previous call; the peers write it again two calls later. The release arrive + // orders the reads of the flags before CTA 0 rotates them. + uint4 const flags = reinterpret_cast(p.exchangeFlags)[0]; + uint4 const sizes = reinterpret_cast(p.exchangeFlags)[1]; + clusterArrive(); + if (flags.w > 0) + { + auto* dirty = reinterpret_cast( + static_cast(p.exchangeLocal) + static_cast(flags.y) * flags.z); + int const dirtyPackets = static_cast(sizes.x / kDraftSampleGluePacketBytes); + for (int i = lane; i < dirtyPackets; i += 32) + { + dirty[i] = make_uint4(kEmptyWord, kEmptyWord, kEmptyWord, kEmptyWord); + } + } + } + else + { + clusterArriveRelaxed(); + } + } + else if (warp == kWideStateWarp) + { + // The next step's state does not need the tokens; nothing in the cluster waits for it. + clusterArriveRelaxed(); + if (cta == 0) + { + advanceStep(p); + } + } + else + { + // The hnorm group: hnorm of hidden row cta does not need the tokens; nothing in the cluster waits for it. + clusterArriveRelaxed(); + if (hiddenHalf && cta < p.numRows) + { + int const t = tid % kNormThreads; + auto const* src = reinterpret_cast(p.hidden + cta * p.hiddenStride); +#pragma unroll + for (int b = 0; b < kNormVecs; ++b) + { + cpAsync16(&sHidden[t + kNormThreads * b], &src[t + kNormThreads * b]); + } + // Each thread reads back only the vectors it copied. + cpAsyncWaitAll(); + int const lo = max(p.colLo - kDraftSampleGlueHidden, 0); + int const hi = min(p.colHi - kDraftSampleGlueHidden, kDraftSampleGlueHidden); + int const outShift = kDraftSampleGlueHidden + lo - p.colLo; + normRowWindow(sHidden, sWeights[1], p.eps, lo, hi, p.glueOut + cta * p.glueOutStride + outShift, + kWideHnormGroup, sGroupSums); + } + } + + // CTA 0 outlives the pushes into its shared memory. + if (!waited) + { + clusterWait(); + } +#endif +} + } // namespace void invokeDraftSampleGlue(DraftSampleGlueParams const& params, cudaStream_t stream) @@ -951,6 +1295,26 @@ void invokeDraftSampleGlue(DraftSampleGlueParams const& params, cudaStream_t str "draft sample glue: the window needs the hidden rows"); TLLM_CHECK_WITH_INFO(params.reqIdxPerToken == nullptr || (params.seqLens != nullptr && params.kvLens != nullptr), "draft sample glue: the DSA metadata needs the sequence and KV lengths"); + bool const wide = params.tpSize == 4; + if (wide) + { + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kWideCtas; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL() ? 1 : 0; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(kWideCtas); + config.blockDim = dim3(kWideThreads); + config.dynamicSmemBytes = 0; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = 2; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, draftSampleGlueWideKernel, params)); + return; + } tensorrt_llm::common::launchWithPdlWhenEnabled( "draftSampleGlue", draftSampleGlueKernel, dim3(1), dim3(kBlockSize), 0, stream, params); } diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.h b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmDraftSampleGlue.h similarity index 100% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.h rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmDraftSampleGlue.h diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.cu b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmDsaHistTopK.cu similarity index 99% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.cu rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmDsaHistTopK.cu index a73b31395b79..6044a1e19409 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.cu +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmDsaHistTopK.cu @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmDsaHistTopK.h" #include "tensorrt_llm/common/envUtils.h" diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.h b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmDsaHistTopK.h similarity index 100% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.h rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmDsaHistTopK.h diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.cu b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmFusedAGemm.cu similarity index 99% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.cu rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmFusedAGemm.cu index 5224fcb38651..f430c16617ff 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.cu +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmFusedAGemm.cu @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmFusedAGemm.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaDriverWrapper.h" diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.h b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmFusedAGemm.h similarity index 100% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.h rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmFusedAGemm.h diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.cu b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmFusedEmbeddingLookup.cu similarity index 97% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.cu rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmFusedEmbeddingLookup.cu index 282fd4fe89a7..551bc5dc9a5b 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.cu +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmFusedEmbeddingLookup.cu @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmFusedEmbeddingLookup.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/envUtils.h" diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.h b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmFusedEmbeddingLookup.h similarity index 100% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.h rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmFusedEmbeddingLookup.h diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmIntraNodeLamport.cuh b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmIntraNodeLamport.cuh new file mode 100644 index 000000000000..2236a36dde8d --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmIntraNodeLamport.cuh @@ -0,0 +1,212 @@ +/* + * 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. + */ +#pragma once + +// Shared-state helpers for the owned fast variants of the in-node one-shot Lamport kernels +// (allreduce_fusion_kernel_oneshot_lamport_fast in allReduceFusionKernels.cu, +// moefinalize_allreduce_fusion_kernel_oneshot_lamport_fast in moeAllReduceFusionKernels.cu). They use the workspace of +// allocate_allreduce_fusion_workspace exactly as the default kernels do: the same buffers, flag rotation, counter and +// sizes to clear, so either kind of kernel may follow the other on one workspace. + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::ipc +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + +// Brings the line of ptr into L2 without reading it. Safe before a grid dependency wait for any address: no value +// reaches a register, and a later write of the line (L2 is the point of coherence) updates the cached copy. +__device__ __forceinline__ void prefetchL2(void const* ptr) +{ + asm volatile("prefetch.global.L2 [%0];" ::"l"(ptr)); +} + +// The Lamport state of the one-shot kernels, read in two steps. The workspace pointers and the Lamport buffer size are +// written by the host when the workspace is created and by no kernel, so they are read before the grid dependency +// wait, and thread 0 brings the lines of the flag and of the size to clear into L2. Those two are written by the +// previous one-shot kernel on the workspace and are read after the wait (loadFlags). arrive() then derives the buffers +// and counts this CTA's arrival, as LamportComm's constructor does. +template +struct LamportCommSplit +{ + __device__ __forceinline__ LamportCommSplit(void* const* workspace, int rank, int32_t* flags, int64_t* layout) + { + counter_ptr = &flags[0]; + flag_ptr = &flags[2]; + clear_ptr = &layout[0]; + if (threadIdx.x == 0) + { + prefetchL2(flag_ptr); + prefetchL2(clear_ptr); + } + comm_size = layout[1]; +#pragma unroll + for (int r = 0; r < NRanks; ++r) + { + bases[r] = reinterpret_cast(workspace[2 * NRanks + r]); + } + local_base = reinterpret_cast(workspace[2 * NRanks + rank]); + } + + __device__ __forceinline__ void loadFlags() + { + flag_value = *flag_ptr; + clear_size = *clear_ptr; + } + + __device__ __forceinline__ void arrive() + { + int data_offset = flag_value % 3; + int clear_offset = (flag_value + 2) % 3; +#pragma unroll + for (int r = 0; r < NRanks; ++r) + { + data_bufs[r] = bases[r] + data_offset * comm_size; + } + local_data = local_base + data_offset * comm_size; + clear_buf = local_base + clear_offset * comm_size; + __syncthreads(); + if (threadIdx.x == 0) + { + atomicAdd(counter_ptr, 1); + } + } + + __device__ __forceinline__ int loadArrivals() const + { + return *reinterpret_cast(counter_ptr); + } + + // The rotation of LamportComm::update, by CTA 0's thread 0: once the count of arrived CTAs (arrivals, the last + // value read) covers the grid, the next flag, the size the next kernel clears and a reset counter. + __device__ __forceinline__ void rotate(int arrivals, int64_t new_clear_size) + { + while (arrivals != static_cast(gridDim.x)) + { + arrivals = loadArrivals(); + } + *flag_ptr = (flag_value + 1) % 3; + *clear_ptr = new_clear_size; + *counter_ptr = 0; + } + + int* counter_ptr; + int* flag_ptr; + int64_t* clear_ptr; + int64_t comm_size; + uint8_t* bases[NRanks]; + uint8_t* local_base; + uint8_t* data_bufs[NRanks]; + // data_bufs[rank], without a runtime index into the array. + uint8_t* local_data; + uint8_t* clear_buf; + int64_t clear_size; + int flag_value; +}; + +int constexpr kNormExchangeMaxBlocks = 8; + +// The RMSNorm partials of a cluster that covers one row: one float per CTA, and the transaction barrier that completes +// once every CTA's partial has landed. +struct NormExchangeStorage +{ + float partials[kNormExchangeMaxBlocks]; + uint64_t barrier; +}; + +__device__ __forceinline__ NormExchangeStorage& normExchangeStorage() +{ + __shared__ NormExchangeStorage storage; + return storage; +} + +// Before the grid dependency wait, by every thread of a cluster of 2 to kNormExchangeMaxBlocks CTAs: thread 0 sets up +// this CTA's barrier for one arrival (its own) and one partial from every CTA, then every thread arrives on the cluster +// barrier. That phase publishes the barrier before any CTA sends (exchangeNormPartials waits for it). +__device__ __forceinline__ void initNormExchange() +{ +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + if (threadIdx.x == 0) + { + uint32_t const bytes = static_cast(cooperative_groups::this_cluster().num_blocks() * sizeof(float)); + uint32_t const barrier = static_cast(__cvta_generic_to_shared(&normExchangeStorage().barrier)); + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(barrier) : "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(barrier), "r"(bytes) : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + asm volatile("barrier.cluster.arrive.relaxed;" ::: "memory"); +#endif +} + +// Called by every thread after initNormExchange; blockSum is this CTA's partial (thread 0's value is used). Thread 0 +// stores it into slot [block rank] of every CTA of the cluster with st.async, which completes that CTA's barrier, then +// waits on its own barrier and returns the partials added in block-rank order starting from zero: the value the +// cluster.sync() loop of the default epilogue computes. Other threads return 0. +__device__ __forceinline__ float exchangeNormPartials(float blockSum) +{ + float sum = 0.F; +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("barrier.cluster.wait;" ::: "memory"); + if (threadIdx.x == 0) + { + cooperative_groups::cluster_group cluster = cooperative_groups::this_cluster(); + int const numBlocks = static_cast(cluster.num_blocks()); + uint32_t const blockRank = cluster.block_rank(); + auto& storage = normExchangeStorage(); + uint32_t const slot = static_cast(__cvta_generic_to_shared(&storage.partials[blockRank])); + uint32_t const barrier = static_cast(__cvta_generic_to_shared(&storage.barrier)); + for (int block = 0; block < numBlocks; ++block) + { + uint32_t peerSlot; + uint32_t peerBarrier; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(peerSlot) : "r"(slot), "r"(block)); + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(peerBarrier) : "r"(barrier), "r"(block)); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"(peerSlot), + "r"(__float_as_uint(blockSum)), "r"(peerBarrier) + : "memory"); + } + uint32_t complete = 0; + while (!complete) + { + asm volatile( + "{\n" + ".reg .pred done;\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 done, [%1], 0;\n" + "selp.u32 %0, 1, 0, done;\n" + "}\n" + : "=r"(complete) + : "r"(barrier) + : "memory"); + } + for (int block = 0; block < numBlocks; ++block) + { + sum += storage.partials[block]; + } + } +#endif + return sum; +} + +#endif +} // namespace kernels::glm_moe_dsa::ipc + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmIntraNodeReduce.cu b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmIntraNodeReduce.cu new file mode 100644 index 000000000000..8fc4477d3eeb --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmIntraNodeReduce.cu @@ -0,0 +1,600 @@ +/* + * Copyright (c) 2022-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/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/common/reduceKernelUtils.cuh" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmIntraNodeLamport.cuh" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmIntraNodeReduce.h" +#include +#include +#include +#include +#include +#include +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ +namespace +{ +using DType = __nv_bfloat16; +namespace cg = cooperative_groups; +constexpr int kFinalizeFastMaxTopK = 8; +constexpr int kFinalizeFastMaxThreads = 768; +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 +template +__device__ __forceinline__ PackedType add128(PackedType const& a, PackedType const& b) +{ + static constexpr int kMathCount = sizeof(PackedType) / sizeof(DType); + PackedType c; +#pragma unroll + for (int i = 0; i < kMathCount; ++i) + { + reinterpret_cast(&c)[i] = reinterpret_cast(&a)[i] + reinterpret_cast(&b)[i]; + } + return c; +} + +// FastNorm: a cluster's partials are exchanged with ipc::exchangeNormPartials (the kernel called +// ipc::initNormExchange before its grid dependency wait) instead of two cluster barriers; same sum. +template +__device__ __forceinline__ PackedType rms_norm( + PackedType const& residual, PackedType const& gamma, float const eps, int hidden_dim) +{ + static constexpr int kMathCount = sizeof(PackedType) / sizeof(DType); + __shared__ float s_val; + PackedType norm_out; + cg::cluster_group cluster = cg::this_cluster(); + float acc = 0.f; +#pragma unroll + for (int i = 0; i < kMathCount; ++i) + { + float v = static_cast(reinterpret_cast(&residual)[i]); + acc += v * v; + } + tensorrt_llm::common::blockReduceSumV2(&acc); + if (cluster.num_blocks() > 1) + { + if constexpr (FastNorm) + { + float const clusterSum = ipc::exchangeNormPartials(acc); + if (threadIdx.x == 0) + { + acc = clusterSum; + } + } + else + { + if (threadIdx.x == 0) + { + s_val = acc; + acc = 0.f; + } + cluster.sync(); + if (threadIdx.x == 0) + { + for (int i = 0; i < cluster.num_blocks(); ++i) + { + acc += *cluster.map_shared_rank(&s_val, i); + } + } + cluster.sync(); + } + } + if (threadIdx.x == 0) + { + s_val = rsqrtf(acc / hidden_dim + eps); + } + __syncthreads(); +#pragma unroll + for (int i = 0; i < kMathCount; ++i) + { + reinterpret_cast(&norm_out)[i] + = static_cast(static_cast(reinterpret_cast(&residual)[i]) * s_val + * static_cast(reinterpret_cast(&gamma)[i])); + } + return norm_out; +} + +__device__ __forceinline__ bool is_neg_zero(float v) +{ + return *reinterpret_cast(&v) == 0x80000000; +} + +__device__ __forceinline__ bool is_neg_zero(float4 v) +{ + return is_neg_zero(v.x) || is_neg_zero(v.y) || is_neg_zero(v.z) || is_neg_zero(v.w); +} + +__device__ __forceinline__ float4 get_neg_zero() +{ + float4 vec; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + reinterpret_cast(&vec)[i] = 0x80000000; + } + return vec; +} + +__device__ __forceinline__ float4 ld_global_volatile(float4* addr) +{ + float4 val; + asm volatile("ld.volatile.global.v4.f32 {%0, %1, %2, %3}, [%4];" + : "=f"(val.x), "=f"(val.y), "=f"(val.z), "=f"(val.w) + : "l"(addr)); + return val; +} + +template +__device__ __forceinline__ float4 allreduce_sum(float4* vals) +{ + if constexpr (Fp32Acc) + { + static_assert(!std::is_same_v); + float acc_f32[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) + { + acc_f32[i] = static_cast(reinterpret_cast(&vals[0])[i]); + } +#pragma unroll + for (int r = 1; r < NRanks; ++r) + { +#pragma unroll + for (int i = 0; i < 8; ++i) + { + acc_f32[i] += static_cast(reinterpret_cast(&vals[r])[i]); + } + } + float4 acc; +#pragma unroll + for (int i = 0; i < 8; ++i) + { + reinterpret_cast(&acc)[i] = static_cast(acc_f32[i]); + } + return acc; + } + else + { + float4 acc = vals[0]; +#pragma unroll + for (int r = 1; r < NRanks; ++r) + { + acc = add128(acc, vals[r]); + } + return acc; + } +} + +#endif +template +__global__ void __launch_bounds__(1024) reduceKernel(IntraNodeReduceParams params) +{ +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + auto cluster = cg::this_cluster(); + auto grid = cg::this_grid(); + int token_id = grid.cluster_rank(); + int access_id_in_token = cluster.thread_rank(); + int access_id = token_id * params.hidden_dim / 8 + access_id_in_token; + int access_stride = grid.num_clusters() * params.hidden_dim / 8; + int tot_access = params.size / 8; + float4 clear_vec = get_neg_zero(); + float4 const gamma = reinterpret_cast(params.rms_gamma)[access_id_in_token]; + float4 const residual = reinterpret_cast(params.residual_in)[access_id]; + ipc::LamportCommSplit comm(params.workspace, params.rank, params.flags, params.layout); + if (cg::this_cluster().num_blocks() > 1) + { + ipc::initNormExchange(); + } + + cudaGridDependencySynchronize(); + if constexpr (!TriggerCompletionAtEnd) + { + cudaTriggerProgrammaticLaunchCompletion(); + } + alignas(16) float val[4]; + *reinterpret_cast(val) = reinterpret_cast(params.allreduce_in)[access_id]; + comm.loadFlags(); + comm.arrive(); + int clear_access = comm.clear_size / 8; + +#pragma unroll + for (int i = 0; i < 4; ++i) + { + if (is_neg_zero(val[i])) + { + val[i] = 0.f; + } + } + float4 const own = *reinterpret_cast(val); +#pragma unroll + for (int r = 0; r < NRanks; ++r) + { + reinterpret_cast(comm.data_bufs[r])[params.rank * tot_access + access_id] = own; + } + for (int idx = access_id; idx < clear_access; idx += access_stride) + { + reinterpret_cast(comm.clear_buf)[idx] = clear_vec; + } + + bool rotating = blockIdx.x == 0 && threadIdx.x == 0; + int arrivals = rotating ? comm.loadArrivals() : 0; + float4 vals[NRanks]; + bool done = false; + while (!done) + { + done = true; +#pragma unroll + for (int r = 0; r < NRanks; ++r) + { + if (r != params.rank) + { + vals[r] = ld_global_volatile(&reinterpret_cast(comm.local_data)[r * tot_access + access_id]); + done &= !is_neg_zero(vals[r]); + } + } + } +#pragma unroll + for (int r = 0; r < NRanks; ++r) + { + if (r == params.rank) + { + vals[r] = own; + } + } + if (rotating && arrivals == static_cast(gridDim.x)) + { + comm.rotate(arrivals, params.size * NRanks); + rotating = false; + } + float4 sum_val = allreduce_sum(vals); + float4 const added = add128(sum_val, residual); + reinterpret_cast(params.residual_out)[access_id] = added; + reinterpret_cast(params.norm_out)[access_id] + = rms_norm(added, gamma, params.rms_eps, params.hidden_dim); + if (rotating) + { + comm.rotate(arrivals, params.size * NRanks); + } + + if constexpr (TriggerCompletionAtEnd) + { + cudaTriggerProgrammaticLaunchCompletion(); + } +#endif +} + +template +__global__ void __launch_bounds__(kFinalizeFastMaxThreads) finalizeKernel(IntraNodeReduceParams params) +{ +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + cg::grid_group grid = cg::this_grid(); + + int token_id = grid.cluster_rank(); + int access_id_in_token = cluster.thread_rank(); + int access_id = token_id * params.hidden_dim / 8 + access_id_in_token; + int access_stride = grid.num_clusters() * params.hidden_dim / 8; + int tot_access = params.size / 8; + int thread_offset_within_token = access_id_in_token * 8; + float4 clear_vec = get_neg_zero(); + + union ACC_TYPE + { + float4 packed; + DType unpacked[8]; + }; + + float4 const gamma_val = reinterpret_cast(params.rms_gamma)[access_id_in_token]; + ipc::LamportCommSplit comm(params.workspace, params.rank, params.flags, params.layout); + if (cluster.num_blocks() > 1) + { + ipc::initNormExchange(); + } + // The lines read right after the wait, brought into L2 while it is pending: the token's routing (thread 0), and + // the residual and shared-expert rows (one thread per 128-byte line). + if (threadIdx.x == 0) + { + ipc::prefetchL2(params.expanded_idx_to_permuted_idx + token_id * params.top_k); + if (params.expert_scale_factor) + { + ipc::prefetchL2(static_cast(params.expert_scale_factor) + token_id * params.top_k); + } + } + if (access_id_in_token % (128 / sizeof(float4)) == 0) + { + if (params.residual_in) + { + ipc::prefetchL2(reinterpret_cast(params.residual_in) + access_id); + } + if (params.shared_expert_output) + { + ipc::prefetchL2(reinterpret_cast(params.shared_expert_output) + access_id); + } + } + + cudaGridDependencySynchronize(); + cudaTriggerProgrammaticLaunchCompletion(); + + int const top_k = params.top_k; + bool const use_scale_factor = params.expert_scale_factor != nullptr; + int32_t permuted_idx[kFinalizeFastMaxTopK]; + float block_scale[kFinalizeFastMaxTopK]; +#pragma unroll + for (int k = 0; k < kFinalizeFastMaxTopK; ++k) + { + permuted_idx[k] = -1; + block_scale[k] = 1.0; + if (k < top_k) + { + int const expanded_idx = token_id * top_k + k; + permuted_idx[k] = params.expanded_idx_to_permuted_idx[expanded_idx]; + if (use_scale_factor) + { + block_scale[k] + = static_cast(static_cast(params.expert_scale_factor)[expanded_idx]); + } + } + } + ACC_TYPE shared_expert_output; + shared_expert_output.packed = make_float4(0.f, 0.f, 0.f, 0.f); + if (params.shared_expert_output) + { + shared_expert_output.packed = reinterpret_cast(params.shared_expert_output)[access_id]; + } + float4 residual_val = make_float4(0.f, 0.f, 0.f, 0.f); + if (params.residual_in) + { + residual_val = reinterpret_cast(params.residual_in)[access_id]; + } + comm.loadFlags(); + + ACC_TYPE permuted_data[kFinalizeFastMaxTopK]; +#pragma unroll + for (int k = 0; k < kFinalizeFastMaxTopK; ++k) + { + if (permuted_idx[k] != -1) + { + int thread_offset_across_token = permuted_idx[k] * params.hidden_dim + thread_offset_within_token; + permuted_data[k].packed + = reinterpret_cast(params.allreduce_in)[thread_offset_across_token / 8]; + } + } + comm.arrive(); + int clear_access = comm.clear_size / 8; + + // The finalize of the default kernel: top-k slots in order, -1 slots skipped, then the shared expert's output, in + // fp32, rounded to DType once. + float facc[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) + { + facc[i] = 0.f; + } +#pragma unroll + for (int k = 0; k < kFinalizeFastMaxTopK; ++k) + { + if (permuted_idx[k] == -1) + continue; +#pragma unroll + for (int i = 0; i < 8; ++i) + { + facc[i] += static_cast(permuted_data[k].unpacked[i]) * block_scale[k]; + } + } + if (params.shared_expert_output) + { +#pragma unroll + for (int i = 0; i < 8; ++i) + { + facc[i] += static_cast(shared_expert_output.unpacked[i]); + } + } + ACC_TYPE accumulator; +#pragma unroll + for (int i = 0; i < 8; ++i) + { + accumulator.unpacked[i] = static_cast(facc[i]); + } + alignas(16) float val[4] = {accumulator.packed.x, accumulator.packed.y, accumulator.packed.z, accumulator.packed.w}; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + if (is_neg_zero(val[i])) + { + val[i] = 0.f; + } + } + float4 const own = *reinterpret_cast(val); +#pragma unroll + for (int r = 0; r < NRanks; ++r) + { + reinterpret_cast(comm.data_bufs[r])[params.rank * tot_access + access_id] = own; + } + for (int idx = access_id; idx < clear_access; idx += access_stride) + { + reinterpret_cast(comm.clear_buf)[idx] = clear_vec; + } + + bool rotating = blockIdx.x == 0 && threadIdx.x == 0; + int arrivals = rotating ? comm.loadArrivals() : 0; + float4 vals[NRanks]; + bool done = false; + while (!done) + { + done = true; +#pragma unroll + for (int r = 0; r < NRanks; ++r) + { + if (r != params.rank) + { + vals[r] = ld_global_volatile(&reinterpret_cast(comm.local_data)[r * tot_access + access_id]); + done &= !is_neg_zero(vals[r]); + } + } + } +#pragma unroll + for (int r = 0; r < NRanks; ++r) + { + if (r == params.rank) + { + vals[r] = own; + } + } + if (rotating && arrivals == static_cast(gridDim.x)) + { + comm.rotate(arrivals, params.size * NRanks); + rotating = false; + } + float4 sum_val = vals[0]; +#pragma unroll + for (int r = 1; r < NRanks; ++r) + { + sum_val = add128(sum_val, vals[r]); + } + + // The epilogue of fused_op, with the residual and gamma already loaded. + float4 norm_input; + if (params.residual_in) + { + norm_input = add128(sum_val, residual_val); + reinterpret_cast(params.residual_out)[access_id] = norm_input; + } + else + { + norm_input = sum_val; + } + float4 norm_val = rms_norm(norm_input, gamma_val, params.rms_eps, params.hidden_dim); + reinterpret_cast(params.norm_out)[access_id] = norm_val; + if (rotating) + { + comm.rotate(arrivals, params.size * NRanks); + } +#endif +} + +struct LaunchShape +{ + int cluster{}; + int block{}; + bool valid{}; +}; + +bool fp32Acc() +{ + static char const* value = std::getenv("ALL_REDUCE_FUSION_KERNEL_ACC_FP32"); + return value == nullptr || value[0] != '0'; +} + +LaunchShape launchShape(int rows, int hidden, int ranks, bool finalize) +{ + if (rows < 1 || rows > 24 || hidden != 6144 || ranks != 4) + return {}; + int device; + TLLM_CUDA_CHECK(cudaGetDevice(&device)); + using Key = std::tuple; + static thread_local std::map cache; + Key key{device, rows, finalize}; + if (auto it = cache.find(key); it != cache.end()) + return it->second; + cudaDeviceProp prop; + TLLM_CUDA_CHECK(cudaGetDeviceProperties(&prop, device)); + LaunchShape shape{}; + if (prop.major == 10 && prop.minor == 7) + { + int cluster = 8; + if (finalize) + { + if (common::getBoolEnv("TRTLLM_MOE_FINALIZE_CTA1")) + { + int active = 0; + TLLM_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&active, finalizeKernel<4>, 768, 0)); + if (active > 0) + cluster = 1; + } + } + else + { + while (768 / cluster < 128 || (rows * cluster > prop.multiProcessorCount && cluster > 1)) + cluster /= 2; + } + int block = 768 / cluster; + int active = 0; + if (finalize) + TLLM_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&active, finalizeKernel<4>, block, 0)); + else if (fp32Acc()) + TLLM_CUDA_CHECK( + cudaOccupancyMaxActiveBlocksPerMultiprocessor(&active, reduceKernel<4, true, true>, block, 0)); + else + TLLM_CUDA_CHECK( + cudaOccupancyMaxActiveBlocksPerMultiprocessor(&active, reduceKernel<4, false, true>, block, 0)); + shape = {cluster, block, active > 0 && rows * cluster <= prop.multiProcessorCount}; + } + cache.emplace(key, shape); + return shape; +} + +template +void launch(Kernel kernel, IntraNodeReduceParams const& p, LaunchShape shape) +{ + cudaLaunchAttribute attrs[2]{}; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = common::getEnvEnablePDL() ? 1 : 0; + attrs[1].id = cudaLaunchAttributeClusterDimension; + attrs[1].val.clusterDim = {static_cast(shape.cluster), 1, 1}; + cudaLaunchConfig_t cfg{}; + cfg.gridDim = (p.size / p.hidden_dim) * shape.cluster; + cfg.blockDim = shape.block; + cfg.stream = p.stream; + cfg.attrs = attrs; + cfg.numAttrs = 2; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&cfg, kernel, p)); +} +} // namespace + +bool intraNodeReduceSupported(int rows, int hidden, int nranks, bool finalize) +{ + return launchShape(rows, hidden, nranks, finalize).valid; +} + +void intraNodeReduce(IntraNodeReduceParams const& p) +{ + auto shape = launchShape(p.size / p.hidden_dim, p.hidden_dim, p.nranks, false); + TLLM_CHECK_WITH_INFO(shape.valid, "Unsupported GLM Rubin IPC reduce launch"); + if (fp32Acc()) + { + if (p.trigger_completion_at_end) + launch(reduceKernel<4, true, true>, p, shape); + else + launch(reduceKernel<4, true, false>, p, shape); + } + else if (p.trigger_completion_at_end) + launch(reduceKernel<4, false, true>, p, shape); + else + launch(reduceKernel<4, false, false>, p, shape); +} + +void intraNodeMoeFinalize(IntraNodeReduceParams const& p) +{ + auto shape = launchShape(p.size / p.hidden_dim, p.hidden_dim, p.nranks, true); + TLLM_CHECK_WITH_INFO(shape.valid && p.top_k >= 1 && p.top_k <= 8, "Unsupported GLM Rubin IPC finalize launch"); + launch(finalizeKernel<4>, p, shape); +} +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmIntraNodeReduce.h b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmIntraNodeReduce.h new file mode 100644 index 000000000000..45651a9abd8a --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmIntraNodeReduce.h @@ -0,0 +1,53 @@ +/* + * Copyright (c) 2022-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/common/config.h" +#include +#include +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ +// IPC pointer table and its real mutable state. This is not the MNNVL ABI. +struct IntraNodeReduceParams +{ + void* const* workspace{}; + int32_t* flags{}; + int64_t* layout{}; + void const* allreduce_in{}; + void const* residual_in{}; + void const* rms_gamma{}; + void* residual_out{}; + void* norm_out{}; + void const* expert_scale_factor{}; + int32_t const* expanded_idx_to_permuted_idx{}; + void const* shared_expert_output{}; + int rank{}; + int nranks{}; + int size{}; + int hidden_dim{}; + int top_k{}; + float rms_eps{}; + bool trigger_completion_at_end{true}; + cudaStream_t stream{}; +}; + +bool intraNodeReduceSupported(int rows, int hidden, int nranks, bool finalize); +void intraNodeReduce(IntraNodeReduceParams const& params); +void intraNodeMoeFinalize(IntraNodeReduceParams const& params); +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLamport.cuh b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmLamport.cuh similarity index 100% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmLamport.cuh rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmLamport.cuh diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmLowMBf16Moe.cu b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmLowMBf16Moe.cu new file mode 100644 index 000000000000..f8c73895af05 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmLowMBf16Moe.cu @@ -0,0 +1,1664 @@ +/* + * 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/glmMoeDsaRubin/glmLowMBf16Moe.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/kernels/moeTopKFuncs.cuh" + +#include +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ +namespace +{ +namespace cg = cooperative_groups; + +using Bf16 = __nv_bfloat16; + +constexpr int kNumExperts = kLowMBf16MoeNumExperts; +constexpr int kTopK = kLowMBf16MoeTopK; +static_assert(kNumExperts % 32 == 0, "routing gives each lane of one warp kNumExperts / 32 candidates"); + +// FC1: 8 consumer warps and one producer warp per CTA. Consumer thread t reads the 16-byte pieces t, +// t + kFc1ConsumerThreads, ... of x and of every weight row, so the hidden size is a multiple of kHiddenGranule and at +// most kMaxPiecesPerThread of them. +constexpr int kFc1ConsumerWarps = 8; +constexpr int kFc1ConsumerThreads = kFc1ConsumerWarps * 32; +constexpr int kFc1Threads = kFc1ConsumerThreads + 32; +constexpr int kPieceElems = 8; +constexpr int kHiddenGranule = kFc1ConsumerThreads * kPieceElems; +constexpr int kMaxPiecesPerThread = 4; +// The FC1 ring holds whole weight rows, as many as fit in kFc1RingBytes and at most kMaxStages. At hidden size 6144 +// and interSize 256 that is 7 rows of 12,288 B, and an FC1 CTA (ring, row sums, barriers and the 1,024 B the runtime +// reserves per CTA) takes 87,840 B, so one FC2 CTA (138,312 B) fits beside it in the 233,472 B of an SM. +constexpr int kFc1RingBytes = 7 * 6144 * static_cast(sizeof(Bf16)); +constexpr int kMaxStages = 8; + +// FC2: a warp per routed expert. +constexpr int kFc2Threads = 256; +constexpr int kFc2Warps = kFc2Threads / 32; +static_assert(kTopK == kFc2Warps, "the FC2 kernel gives each routed expert one warp"); + +// FC2: mma.sync m16n8k16 tiles have 16 matrix rows; one 32-element K step is two mma calls (see mmaStep). A warp per +// routed expert walks the intermediate dimension in pairs of K steps, one per accumulator. +constexpr int kMmaRows = 16; +constexpr int kStepK = 32; +constexpr int kFc2StepPairK = 2 * kStepK; +// The largest intermediate size per rank whose FC2 working set fits in shared memory (kTiles = 1). +constexpr int kMaxInterSize = 512; + +// The routed scaling factor as the FC2 kernel takes it: the double value of noaux_tc_op's FP64 weight expression, its +// float value, and whether routingWeight's FP32 path applies to it (decided on the host, see makeRoutingScale). +struct RoutingScale +{ + double value; + float valueFloat; + int fastPath; +}; + +// Multi-token calls (invokeLowMBf16MoeMulti). The FC1 kernel keeps x of every token in shared memory, of which a +// consumer thread reads its at most kMultiMaxPiecesPerThread pieces per token (hidden size 6144), and forms the SwiGLU +// outputs every kMultiChunkUnits units from per-warp row sums held in two alternating buffers. Its ring has between +// kMultiMinFc1Stages (where the shared memory of the SM allows) and kMultiMaxFc1Stages rows. +constexpr int kMultiMaxTokens = kLowMBf16MoeMaxTokens; +constexpr int kMultiMaxPiecesPerThread = 3; +constexpr int kMultiChunkUnits = 8; +constexpr int kMultiChunkRows = 2 * kMultiChunkUnits; +constexpr int kMultiMinFc1Stages = 7; +constexpr int kMultiMaxFc1Stages = 12; +// The multi-token FC1 kernel's warps: the consumer warps of lowMBf16MoeFc1Kernel, its producer warp, and an epilogue +// warp that forms the SwiGLU outputs of each chunk while the consumer warps go on with the next. +constexpr int kMultiFc1EpilogueWarp = kFc1ConsumerWarps + 1; +constexpr int kMultiFc1Threads = kFc1ConsumerThreads + 64; +// FC2: kMultiFc2Tiles 16-row tiles per CTA, one consumer warp per (union expert, tile) and one producer warp that +// streams the union experts' weight slices through a ring of at most kMultiFc2MaxStages stages. One tile per CTA makes +// CTAs small enough that one fits beside an FC1 CTA and two share an SM once the FC1 CTAs exit. +constexpr int kMultiFc2Tiles = 1; +constexpr int kMultiFc2Rows = kMultiFc2Tiles * kMmaRows; +constexpr int kMultiFc2ConsumerWarps = 8; +constexpr int kMultiFc2ConsumerThreads = kMultiFc2ConsumerWarps * 32; +constexpr int kMultiFc2Threads = kMultiFc2ConsumerThreads + 32; +constexpr int kMultiFc2MaxStages = 8; +// An m16n8k16 mma of the FC2 kernel takes four tokens, two B columns each (see lowMBf16MoeMultiFc2Kernel). +constexpr int kMultiTokensPerMma = 4; +static_assert(kMultiFc2ConsumerWarps % kMultiFc2Tiles == 0, "the consumer warps take whole union experts"); +static_assert(kMultiMaxTokens <= 2 * kMultiTokensPerMma, "the FC2 kernel takes at most two mma token groups"); + +// The routing of a multi-token call as both multi-token kernels hold it in shared memory: the experts each token +// selected, and their union in ascending expert order with, per union expert, the tokens that selected it and its slot +// in each of their lists. +template +struct MultiRouting +{ + // Bit m: token m selected the expert. + uint32_t expertTokens[kNumExperts]; + // Union position of each selected expert. + uint8_t unionPos[kNumExperts]; + int32_t topExperts[kTokens][kTopK]; + int32_t unionExperts[kTokens * kTopK]; + uint32_t unionTokens[kTokens * kTopK]; + // Slot of the union expert in token m's list, where bit m of its unionTokens is set. + uint8_t unionSlots[kTokens * kTopK][kTokens]; + int32_t numUnion; +}; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +constexpr int kCandidates = kNumExperts / 32; +constexpr int kFc1ProducerWarp = kFc1ConsumerWarps; +// Named barrier at which the FC1 consumer warps wait while the producer selects the experts and issues the first rows. +constexpr int kFc1ParkBarrier = 1; +// The FC2 kernel holds its routing computation until this long after its entry; see lowMBf16MoeFc2Kernel. +constexpr uint64_t kFc2RouteHoldNs = 1000; + +using TopKKey = reduce_topk::TopKRedType; + +// Same expression as sigmoid_accurate in noAuxTcKernels.cu. +__device__ __forceinline__ float sigmoidAccurate(float x) +{ + return 0.5f * tanhf(0.5f * x) + 0.5f; +} + +__device__ __forceinline__ uint32_t warpMax(uint32_t value) +{ + uint32_t result; + asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n" : "=r"(result) : "r"(value)); + return result; +} + +// The correction bias of the experts a lane holds (i * 32 + lane) as float. A constant, so the kernels may read it +// before their grid-dependency wait. +template +__device__ __forceinline__ void loadBias(BiasT const* bias, int lane, float (&values)[kCandidates]) +{ +#pragma unroll + for (int i = 0; i < kCandidates; ++i) + { + values[i] = static_cast(bias[i * 32 + lane]); + } +} + +// The routing keys of noaux_tc_op for n_group == 1 (deepseek_v3_topk_kernel, single group) of the experts a lane +// holds: the sigmoid score plus the correction bias, with the expressions of that kernel, packed with the expert index +// as reduceTopK (moeTopKFuncs.cuh) packs them. Stores the sigmoid scores when sigmoidScores is not null. The logits +// are read through L2. +__device__ __forceinline__ void buildKeys(float const* routerLogits, float const (&biasValues)[kCandidates], int lane, + TopKKey (&keys)[kCandidates], float* sigmoidScores) +{ +#pragma unroll + for (int i = 0; i < kCandidates; ++i) + { + int const expert = i * 32 + lane; + float const sigmoid = sigmoidAccurate(__ldcg(routerLogits + expert)); + if (sigmoidScores != nullptr) + { + sigmoidScores[expert] = sigmoid; + } + keys[i] = TopKKey{sigmoid + biasValues[i], expert}; + } +} + +// The top kTopK of one full warp's keys in the order of reduceTopK: descending score, on equal scores the lower index +// first. The per-lane sort is reduceTopK's; each round then takes the warp maximum of the lanes' leading keys with two +// 32-bit reductions, of the score half and then of the index half among the lanes holding the maximum score, which is +// the maximum of reduceTopK's 64-bit reduction. So the experts and their order are noaux_tc_op's. Returns the k-th +// expert on lane k < kTopK, 0 on the other lanes. +__device__ __forceinline__ int32_t selectTopK(TopKKey (&keys)[kCandidates], int lane) +{ + reduce_topk::Sort::run(keys); + + int32_t laneExpert = 0; + uint64_t packedMax = 0; +#pragma unroll + for (int k = 0; k < kTopK; ++k) + { + if (k > 0 && packedMax == keys[0].compValIdx) + { + // This lane held the previous maximum: drop it, as reduceTopK does. +#pragma unroll + for (int i = 0; i + 1 < kCandidates; ++i) + { + keys[i] = keys[i + 1]; + } + keys[kCandidates - 1] = TopKKey{-INFINITY, (kCandidates - 1) * 32 + lane}; + } + uint32_t const scoreBits = static_cast(keys[0].compValIdx >> 32); + uint32_t const maxScoreBits = warpMax(scoreBits); + uint32_t const indexBits = warpMax(scoreBits == maxScoreBits ? static_cast(keys[0].compValIdx) : 0U); + packedMax = (static_cast(maxScoreBits) << 32) | indexBits; + if (lane == k) + { + laneExpert = TopKKey::kMaxIdx - static_cast(indexBits & 0xFFFF); + } + } + return laneExpert; +} + +// noaux_tc_op's routing weight, static_cast(s * c / (S + 1e-20)) in double with s the expert's sigmoid score, +// S the sum of the selected scores and c the routed scaling factor, without FP64 instructions on the common path: FP64 +// in the FC2 CTAs drops the SM clock for the whole FC1 window (measured on Rubin: 2.38 -> 1.6 GHz). When c is a float +// with at most 20 significant bits (fastPath) and S >= 2^-13, the double expression is the float rounding of the exact +// quotient s * c / S: s * c is exact in double, S + 1e-20 rounds to S, and rounding to double first cannot move a +// quotient of these operands across a float rounding boundary except at an exact midpoint, which rounds the same +// either way. The quotient here comes from an exact two-float product, a correctly rounded division and its exact +// remainder; it is returned when it lies more than 2^-20 ulp inside its rounding interval (and is not a power of two, +// where the interval is lopsided). Otherwise, and for scores below 2^-60 (where the FP32 intermediates could lose +// bits), the FP64 expression runs, out of line so that none of it is scheduled on the common path. +__device__ __noinline__ float routingWeightFp64(float s, float S, double c) +{ + return static_cast(s * c / (S + 1e-20)); +} + +__device__ __forceinline__ float routingWeight(float s, float S, RoutingScale const& scale) +{ + if (scale.fastPath != 0 && S >= 0x1p-13f && s >= 0x1p-60f) + { + float const c = scale.valueFloat; + float const pHi = __fmul_rn(s, c); + float const pLo = __fmaf_rn(s, c, -pHi); + float const q = __fdiv_rn(pHi, S); + float const r = __fmaf_rn(-q, S, pHi); + float const e = __fdiv_rn(__fadd_rn(r, pLo), S); + float const y = __fadd_rn(q, e); + uint32_t const bits = __float_as_uint(y); + if (y > 0.F && (bits & 0x7FFFFFU) != 0 && (bits & 0x7F800000U) != 0x7F800000U) + { + float const ulp = __uint_as_float(bits & 0x7F800000U) * 0x1p-23f; + float const g = __fsub_rn(e, __fsub_rn(y, q)); + if (fabsf(g) < 0.5F * ulp - 0x1p-20f * ulp) + { + return y; + } + } + } + return routingWeightFp64(s, S, scale.value); +} + +__device__ __forceinline__ uint32_t smemAddress(void const* ptr) +{ + return static_cast(__cvta_generic_to_shared(ptr)); +} + +__device__ __forceinline__ void mbarrierInit(uint64_t* bar, uint32_t count) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(smemAddress(bar)), "r"(count) : "memory"); +} + +// Makes the barrier initialization visible to the bulk-copy (async) proxy. +__device__ __forceinline__ void fenceBarrierInit() +{ + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); +} + +__device__ __forceinline__ void mbarrierArriveExpectTx(uint64_t* bar, uint32_t bytes) +{ + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(smemAddress(bar)), "r"(bytes) + : "memory"); +} + +// Arrives on the barrier with release semantics: this thread's earlier shared-memory reads are done before the phase +// can complete. +__device__ __forceinline__ void mbarrierArrive(uint64_t* bar) +{ + asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];\n" ::"r"(smemAddress(bar)) : "memory"); +} + +// Whether the barrier's phase with the given parity has completed (one try, which may suspend the thread for a +// while). +__device__ __forceinline__ bool mbarrierTryWait(uint64_t* bar, uint32_t parity) +{ + uint32_t done = 0; + asm volatile( + "{\n" + ".reg .pred complete;\n" + "mbarrier.try_wait.parity.shared::cta.b64 complete, [%1], %2;\n" + "selp.u32 %0, 1, 0, complete;\n" + "}\n" + : "=r"(done) + : "r"(smemAddress(bar)), "r"(parity) + : "memory"); + return done != 0; +} + +__device__ __forceinline__ void mbarrierWait(uint64_t* bar, uint32_t parity) +{ + while (!mbarrierTryWait(bar, parity)) + { + } +} + +// Waits at named barrier `id` until `threads` threads (whole warps) have arrived. +__device__ __forceinline__ void namedBarrierSync(int id, int threads) +{ + asm volatile("bar.sync %0, %1;\n" ::"r"(id), "r"(threads) : "memory"); +} + +__device__ __forceinline__ uint64_t globalTimerNs() +{ + uint64_t now; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(now)); + return now; +} + +// Streamed weights are read once per call: mark their L2 lines evict-first. +__device__ __forceinline__ uint64_t evictFirstPolicy() +{ + uint64_t policy; + asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;\n" : "=l"(policy)); + return policy; +} + +// 1D bulk copy global -> shared that completes `bytes` transactions on `bar`. +__device__ __forceinline__ void bulkCopyToShared( + void* dst, void const* src, uint32_t bytes, uint64_t* bar, uint64_t policy) +{ + uint64_t const srcGlobal = static_cast(__cvta_generic_to_global(src)); + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], %4;\n" + : + : "r"(smemAddress(dst)), "l"(srcGlobal), "r"(bytes), "r"(smemAddress(bar)), "l"(policy) + : "memory"); +} + +__device__ __forceinline__ void mmaBf16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) +{ + asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +// One 32-element K step of a 16-row tile times a vector. Lane (g = lane / 4, t = lane % 4) passes elements [8t, 8t + 8) +// of the step from tile row g (lo), tile row g + 8 (hi) and the vector (vec). The fragment's K columns +// {2t, 2t + 1, 2t + 8, 2t + 9} take the lane's elements {0, 1, 2, 3} in the first call and {4, 5, 6, 7} in the second: +// one K permutation for the matrix and the vector, which leaves the dot products unchanged. Every column of the B +// fragment is the vector, so all accumulator elements hold dot products: elements 0 and 1 of row g, 2 and 3 of row +// g + 8. +__device__ __forceinline__ void mmaStep(float (&acc)[4], uint4 const& lo, uint4 const& hi, uint4 const& vec) +{ + mmaBf16(acc, lo.x, hi.x, lo.y, hi.y, vec.x, vec.y); + mmaBf16(acc, lo.z, hi.z, lo.w, hi.w, vec.z, vec.w); +} + +__device__ __forceinline__ float roundToBf16(float value) +{ + return __bfloat162float(__float2bfloat16_rn(value)); +} + +// The FP32 values of the eight BF16 elements of a 16-byte piece, in memory order. +__device__ __forceinline__ void unpackPiece(uint4 const& piece, float (&values)[kPieceElems]) +{ + uint32_t const words[4] = {piece.x, piece.y, piece.z, piece.w}; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + values[2 * i] = __uint_as_float(words[i] << 16); + values[2 * i + 1] = __uint_as_float(words[i] & 0xFFFF0000U); + } +} + +// Dot product of a 16-byte piece of a BF16 weight row with the FP32 values of the same piece of x, in element order +// with FP32 fused multiply-adds. +__device__ __forceinline__ float dotPiece(uint4 const& piece, float const (&values)[kPieceElems]) +{ + float weights[kPieceElems]; + unpackPiece(piece, weights); + float dot = weights[0] * values[0]; +#pragma unroll + for (int e = 1; e < kPieceElems; ++e) + { + dot = fmaf(weights[e], values[e], dot); + } + return dot; +} + +// The warp sums of two values, each by an xor butterfly over offsets 16, 8, 4, 2, 1, the two shuffle chains +// interleaved. +__device__ __forceinline__ void warpSum2(float& a, float& b) +{ +#pragma unroll + for (int offset = 16; offset > 0; offset /= 2) + { + a += __shfl_xor_sync(0xffffffffU, a, offset); + b += __shfl_xor_sync(0xffffffffU, b, offset); + } +} + +constexpr int kMultiParkBarrier = 1; +constexpr int kMultiRouteBarrier = 2; +constexpr int kMultiConsumerBarrier = 3; + +// Arrives at named barrier `id` of `threads` threads (whole warps) without waiting; the barrier orders this thread's +// earlier shared-memory writes before the waiting threads' later reads. +__device__ __forceinline__ void namedBarrierArrive(int id, int threads) +{ + asm volatile("bar.arrive %0, %1;\n" ::"r"(id), "r"(threads) : "memory"); +} + +// warpSum2's sums with half its shuffles: after the first exchange lanes 0-15 hold a's pairs and lanes 16-31 b's, and +// the next four levels run within each half. Lane 0 ends with warpSum2's a of lane 0 and lane 16 with its b of lane 0: +// the same additions in the same tree (each first pair added in the other operand order, which gives the same sum). +__device__ __forceinline__ float warpSum2Split(float a, float b, int lane) +{ + bool const upper = lane >= 16; + float const received = __shfl_xor_sync(0xffffffffU, upper ? a : b, 16); + float sum = upper ? b + received : a + received; +#pragma unroll + for (int offset = 8; offset > 0; offset /= 2) + { + sum += __shfl_xor_sync(0xffffffffU, sum, offset); + } + return sum; +} + +// dotPiece with x given as a packed piece of eight BF16 values (their FP32 values are exact). +__device__ __forceinline__ float dotPiecePacked(uint4 const& piece, uint4 const& xPiece) +{ + float values[kPieceElems]; + unpackPiece(xPiece, values); + return dotPiece(piece, values); +} + +// One warp builds the union of the tokens' selections (routing.expertTokens and routing.topExperts complete): the +// union experts in ascending expert order, the tokens of each and its slot in each of their lists. +template +__device__ __forceinline__ void buildUnion(MultiRouting& routing, int lane) +{ + constexpr int kPerLane = kNumExperts / 32; + uint32_t tokens[kPerLane]; + int count = 0; +#pragma unroll + for (int i = 0; i < kPerLane; ++i) + { + tokens[i] = routing.expertTokens[lane * kPerLane + i]; + count += tokens[i] != 0 ? 1 : 0; + } + // Inclusive prefix sum of the counts over the lanes. + int prefix = count; +#pragma unroll + for (int offset = 1; offset < 32; offset *= 2) + { + int const value = __shfl_up_sync(0xffffffffU, prefix, offset); + if (lane >= offset) + { + prefix += value; + } + } + int pos = prefix - count; +#pragma unroll + for (int i = 0; i < kPerLane; ++i) + { + if (tokens[i] != 0) + { + int const expert = lane * kPerLane + i; + routing.unionExperts[pos] = expert; + routing.unionTokens[pos] = tokens[i]; + routing.unionPos[expert] = static_cast(pos); + ++pos; + } + } + if (lane == 31) + { + routing.numUnion = prefix; + } + __syncwarp(); + for (int p = lane; p < kTokens * kTopK; p += 32) + { + int const token = p / kTopK; + int const slot = p - token * kTopK; + routing.unionSlots[routing.unionPos[routing.topExperts[token][slot]]][token] = static_cast(slot); + } + __syncwarp(); +} +#endif + +// FC1 + SwiGLU, one CTA per SM (the grid is the SM count). The work is kTopK * interSize units; unit k * interSize + j +// is intermediate index j of routed slot k, that is the expert's gate row (W1, w3w1 row interSize + j) and up row (W3, +// row j). CTA b takes units [b * U / G, (b + 1) * U / G) and streams their rows, gate then up, each as one bulk copy +// into a ring of numStages row buffers. The producer warp selects the experts and issues the rows; it refills a buffer +// as soon as every consumer warp has read its pieces of it, with no CTA-wide barrier per row. The consumer warps wait +// at a named barrier while the producer selects the experts and issues the first rows. Each consumer thread keeps its +// pieces of x in registers and adds its shares of the dot products of two rows at a time, so that the latency chains of +// the two overlap; each consumer warp reduces its shares with shuffles into shared memory, and each unit's SwiGLU adds +// the 8 warp sums in a fixed order. +template +__global__ void __launch_bounds__(kFc1Threads, 2) + lowMBf16MoeFc1Kernel(Bf16* __restrict__ act, Bf16 const* __restrict__ x, float const* __restrict__ routerLogits, + BiasT const* __restrict__ bias, Bf16 const* __restrict__ w3w1, int hiddenSize, int interSize, int numStages) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + extern __shared__ __align__(128) uint8_t smem[]; + // rowFull[s]: buffer s holds its next row (the producer's arrival plus the copy's bytes). rowFree[s]: every + // consumer warp has read its pieces of buffer s. + __shared__ uint64_t rowFull[kMaxStages]; + __shared__ uint64_t rowFree[kMaxStages]; + __shared__ int32_t topExperts[kTopK]; + + int const warp = threadIdx.x / 32; + int const lane = threadIdx.x % 32; + int const rowBytes = hiddenSize * static_cast(sizeof(Bf16)); + int const numUnits = kTopK * interSize; + int const unitBegin = static_cast(static_cast(blockIdx.x) * numUnits / gridDim.x); + int const unitEnd = static_cast(static_cast(blockIdx.x + 1) * numUnits / gridDim.x); + int const numRows = 2 * (unitEnd - unitBegin); + uint8_t* ring = smem; + float* rowSums = reinterpret_cast(smem + numStages * rowBytes); + + if (threadIdx.x == 0) + { + for (int s = 0; s < numStages; ++s) + { + mbarrierInit(&rowFull[s], 1); + mbarrierInit(&rowFree[s], kFc1ConsumerWarps); + } + fenceBarrierInit(); + } + // The correction bias is a constant: the producer reads it before the grid-dependency wait. + float biasValues[kCandidates] = {}; + if (warp == kFc1ProducerWarp) + { + loadBias(bias, lane, biasValues); + } + __syncthreads(); + + // x and the router logits are written by earlier kernels. + cudaGridDependencySynchronize(); + // The FC2 kernel reads only the logits, the bias and its weights before its own wait, and the logits are complete + // since the wait above returned. Released here, its CTAs take the room beside this kernel's and stage their + // weights while this kernel streams. + cudaTriggerProgrammaticLaunchCompletion(); + + if (warp == kFc1ProducerWarp) + { + TopKKey keys[kCandidates]; + buildKeys(routerLogits, biasValues, lane, keys, nullptr); + int32_t const expert = selectTopK(keys, lane); + if (lane < kTopK) + { + topExperts[lane] = expert; + } + __syncwarp(); + + // Lane 0 streams the rows. The next row is of unit (slot, index), the gate row first, into buffer s, which it + // fills for the pass-th time. + uint64_t const policy = evictFirstPolicy(); + int64_t const expertStride = 2 * static_cast(interSize) * hiddenSize; + int slot = unitBegin / interSize; + int index = unitBegin % interSize; + bool gateRow = true; + Bf16 const* expertRows = w3w1 + topExperts[slot] * expertStride; + int s = 0; + uint32_t pass = 0; + auto issueRow = [&]() + { + int64_t const rowInExpert = gateRow ? interSize + index : index; + mbarrierArriveExpectTx(&rowFull[s], rowBytes); + bulkCopyToShared(ring + s * rowBytes, expertRows + rowInExpert * hiddenSize, rowBytes, &rowFull[s], policy); + if (!gateRow && ++index == interSize) + { + // The up row closed the slot's last unit. + index = 0; + if (++slot < kTopK) + { + expertRows = w3w1 + topExperts[slot] * expertStride; + } + } + gateRow = !gateRow; + if (++s == numStages) + { + s = 0; + ++pass; + } + }; + int const firstRows = numStages < numRows ? numStages : numRows; + if (lane == 0) + { + for (int row = 0; row < firstRows; ++row) + { + issueRow(); + } + } + // Release the consumer warps. + __syncwarp(); + namedBarrierSync(kFc1ParkBarrier, kFc1Threads); + if (lane == 0) + { + for (int row = firstRows; row < numRows; ++row) + { + // Buffer s still holds row - numStages until every consumer warp has read it. + mbarrierWait(&rowFree[s], (pass - 1) & 1U); + issueRow(); + } + } + } + else + { + float xValues[kMaxPiecesPerThread][kPieceElems] = {}; + int const piecesPerThread = hiddenSize / kHiddenGranule; +#pragma unroll + for (int i = 0; i < kMaxPiecesPerThread; ++i) + { + if (i < piecesPerThread) + { + unpackPiece( + __ldcg(reinterpret_cast(x) + threadIdx.x + i * kFc1ConsumerThreads), xValues[i]); + } + } + // Wait at the barrier, without polling the ring, while the producer selects the experts and issues the first + // rows. + namedBarrierSync(kFc1ParkBarrier, kFc1Threads); + int s0 = 0; + uint32_t pass0 = 0; + // Rows row and row + 1 (numRows is even) are in buffers s0 and s1 (numStages >= 2), filled for the pass0-th and + // pass1-th time. + for (int row = 0; row < numRows; row += 2) + { + int s1 = s0 + 1; + uint32_t pass1 = pass0; + if (s1 == numStages) + { + s1 = 0; + ++pass1; + } + mbarrierWait(&rowFull[s0], pass0 & 1U); + mbarrierWait(&rowFull[s1], pass1 & 1U); + uint4 const* pieces0 = reinterpret_cast(ring + s0 * rowBytes); + uint4 const* pieces1 = reinterpret_cast(ring + s1 * rowBytes); + float dot0 = dotPiece(pieces0[threadIdx.x], xValues[0]); + float dot1 = dotPiece(pieces1[threadIdx.x], xValues[0]); +#pragma unroll + for (int i = 1; i < kMaxPiecesPerThread; ++i) + { + if (i < piecesPerThread) + { + dot0 += dotPiece(pieces0[threadIdx.x + i * kFc1ConsumerThreads], xValues[i]); + dot1 += dotPiece(pieces1[threadIdx.x + i * kFc1ConsumerThreads], xValues[i]); + } + } + // This warp is done with both buffers (the arrival releases its lanes' reads, ordered before it by the + // warp barrier); the producer refills a buffer once every consumer warp is. + __syncwarp(); + if (lane == 0) + { + mbarrierArrive(&rowFree[s0]); + mbarrierArrive(&rowFree[s1]); + } + warpSum2(dot0, dot1); + if (lane == 0) + { + rowSums[row * kFc1ConsumerWarps + warp] = dot0; + rowSums[(row + 1) * kFc1ConsumerWarps + warp] = dot1; + } + s0 = s1 + 1; + pass0 = pass1; + if (s0 == numStages) + { + s0 = 0; + ++pass0; + } + } + } + __syncthreads(); + + for (int unit = static_cast(threadIdx.x); unit < numRows / 2; unit += kFc1Threads) + { + float gate = 0.F; + float up = 0.F; + for (int w = 0; w < kFc1ConsumerWarps; ++w) + { + gate += rowSums[2 * unit * kFc1ConsumerWarps + w]; + up += rowSums[(2 * unit + 1) * kFc1ConsumerWarps + w]; + } + // The CUTLASS FC1 GEMM writes its output in BF16 and the activation kernel reads it back. + gate = roundToBf16(gate); + up = roundToBf16(up); + // silu(gate) * up in FP32 as the CUTLASS activation (GLUAdaptor) computes it, then BF16. + float const silu = gate * (1.F / (1.F + expf(-gate))); + act[unitBegin + unit] = __float2bfloat16_rn(silu * up); + } +#endif +} + +// FC2 + the weighted sum over the routed experts. Grid: one CTA per kTiles * 16 output rows; warp k computes routed +// expert k's terms from its weight slice, staged in shared memory before the grid-dependency wait. Warp 0 routes before +// that wait, from the router logits and the bias staged by bulk copies, without FP64 instructions. After the wait every +// warp stages its expert's SwiGLU outputs itself. +template +__global__ void __launch_bounds__(kFc2Threads, 2) + lowMBf16MoeFc2Kernel(Bf16* __restrict__ out, Bf16 const* __restrict__ act, float const* __restrict__ routerLogits, + BiasT const* __restrict__ bias, Bf16 const* __restrict__ w2, int hiddenSize, int interSize, RoutingScale scale) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + constexpr int kRows = kTiles * kMmaRows; + extern __shared__ __align__(128) uint8_t smem[]; + __shared__ __align__(16) float logitsShared[kNumExperts]; + __shared__ __align__(16) BiasT biasShared[kNumExperts]; + __shared__ uint64_t inputsFull; + __shared__ uint64_t slicesFull; + __shared__ float sigmoidScores[kNumExperts]; + __shared__ int32_t topExperts[kTopK]; + __shared__ float topWeights[kTopK]; + __shared__ float terms[kTopK][kRows]; + + uint64_t const entryNs = globalTimerNs(); + + int const sliceElems = kRows * interSize; + Bf16* slices = reinterpret_cast(smem); + Bf16* actShared = slices + kTopK * sliceElems; + + int const warp = threadIdx.x / 32; + int const lane = threadIdx.x % 32; + int const row0 = blockIdx.x * kRows; + + if (threadIdx.x == 0) + { + mbarrierInit(&inputsFull, 1); + mbarrierInit(&slicesFull, 1); + fenceBarrierInit(); + } + __syncthreads(); + + if (warp == 0) + { + // Before the grid-dependency wait this kernel reads only the router logits, the bias and the FC2 weights. The + // logits come from the router GEMM, which completed before the FC1 kernel returned from its own wait and + // released this grid; the bias and the weights are constants. The logits and the bias come in by bulk copies: + // plain loads issued now queue behind the FC1 kernel's bulk copies on this SM for microseconds (measured). + bool const bulkInputs + = ((reinterpret_cast(routerLogits) | reinterpret_cast(bias)) % 16) == 0; + if (bulkInputs) + { + if (lane == 0) + { + uint64_t const policy = evictFirstPolicy(); + mbarrierArriveExpectTx(&inputsFull, sizeof(logitsShared) + sizeof(biasShared)); + bulkCopyToShared(logitsShared, routerLogits, sizeof(logitsShared), &inputsFull, policy); + bulkCopyToShared(biasShared, bias, sizeof(biasShared), &inputsFull, policy); + } + mbarrierWait(&inputsFull, 0); + } + else + { + // Each lane stages, and below reads back, only the experts it holds. +#pragma unroll + for (int i = 0; i < kCandidates; ++i) + { + logitsShared[i * 32 + lane] = __ldcg(routerLogits + i * 32 + lane); + biasShared[i * 32 + lane] = bias[i * 32 + lane]; + } + } + // The FC1 kernel's producer warps select the same experts on these SMs at about this time, on the FC1 + // kernel's critical path, and routing here concurrently makes their selection about twice as long (measured + // 0.5 -> 1.1 us). So the routing waits until kFc2RouteHoldNs after this CTA's entry; the weight copies are + // still issued long before the FC1 kernel ends. + while (globalTimerNs() < entryNs + kFc2RouteHoldNs) + { + __nanosleep(64); + } + // The FC1 kernel's routing (the keys of buildKeys, selectTopK) on the staged inputs, and the weights from + // routingWeight: the experts, their order and the FP32 weights of noaux_tc_op. + TopKKey keys[kCandidates]; +#pragma unroll + for (int i = 0; i < kCandidates; ++i) + { + int const expert = i * 32 + lane; + float const sigmoid = sigmoidAccurate(logitsShared[expert]); + sigmoidScores[expert] = sigmoid; + keys[i] = TopKKey{sigmoid + static_cast(biasShared[expert]), expert}; + } + int32_t const expert = selectTopK(keys, lane); + __syncwarp(); + auto const tile = cg::tiled_partition<32>(cg::this_thread_block()); + float const scoreNorm = lane < kTopK ? sigmoidScores[expert] : 0.F; + float const sumNorm = cg::reduce(tile, scoreNorm, cg::plus{}); + if (lane < kTopK) + { + topExperts[lane] = expert; + topWeights[lane] = routingWeight(scoreNorm, sumNorm, scale); + } + __syncwarp(); + if (lane == 0) + { + mbarrierArriveExpectTx(&slicesFull, kTopK * sliceElems * static_cast(sizeof(Bf16))); + } + __syncwarp(); + if (lane < kTopK) + { + bulkCopyToShared(slices + lane * sliceElems, + w2 + (static_cast(topExperts[lane]) * hiddenSize + row0) * interSize, + sliceElems * static_cast(sizeof(Bf16)), &slicesFull, evictFirstPolicy()); + } + } + __syncthreads(); + + // The SwiGLU outputs are written by the FC1 kernel: every warp waits for it, then stages its expert's outputs. + cudaGridDependencySynchronize(); + Bf16 const* slice = slices + warp * sliceElems; + Bf16* expertAct = actShared + warp * interSize; + for (int i = lane; i < interSize / 8; i += 32) + { + reinterpret_cast(expertAct)[i] = __ldcg(reinterpret_cast(act + warp * interSize) + i); + } + __syncwarp(); + mbarrierWait(&slicesFull, 0); + + // Lane (g, t) of warp k accumulates the dot products of slice rows tile * 16 + g and tile * 16 + g + 8 with + // act[k, :], over the two 32-element halves of each 64-element K step in two accumulators. Lanes of odd g take the + // two halves in the other order, so that a warp's loads of 8 rows spread over all 32 banks: column n of the B + // fragments then holds the vector in the order of the rows of n's parity, and the accumulators of those columns + // hold the same products in the same order as with one order for all (which of the two holds which half differs). + int const g = lane / 4; + int const t = lane % 4; + int const swap = g & 1; + float acc[kTiles][2][4] = {}; + for (int j0 = 0; j0 < interSize; j0 += kFc2StepPairK) + { +#pragma unroll + for (int half = 0; half < 2; ++half) + { + int const j = j0 + (half ^ swap) * kStepK + 8 * t; + uint4 const av = *reinterpret_cast(expertAct + j); +#pragma unroll + for (int tileIdx = 0; tileIdx < kTiles; ++tileIdx) + { + uint4 const lo = *reinterpret_cast(slice + (tileIdx * kMmaRows + g) * interSize + j); + uint4 const hi = *reinterpret_cast(slice + (tileIdx * kMmaRows + g + 8) * interSize + j); + mmaStep(acc[tileIdx][half], lo, hi, av); + } + } + } + + // The CUTLASS FC2 epilogue rounds each expert's scaled output to BF16 before adding it into the output. Lane (g, 0) + // holds columns 0 and 1 of rows g and g + 8, of which column g & 1 has their dot products. + float const weight = topWeights[warp]; + if (t == 0) + { +#pragma unroll + for (int tileIdx = 0; tileIdx < kTiles; ++tileIdx) + { + float const lo0 = swap ? acc[tileIdx][0][1] : acc[tileIdx][0][0]; + float const lo1 = swap ? acc[tileIdx][1][1] : acc[tileIdx][1][0]; + float const hi0 = swap ? acc[tileIdx][0][3] : acc[tileIdx][0][2]; + float const hi1 = swap ? acc[tileIdx][1][3] : acc[tileIdx][1][2]; + terms[warp][tileIdx * kMmaRows + g] = roundToBf16(weight * (lo0 + lo1)); + terms[warp][tileIdx * kMmaRows + g + 8] = roundToBf16(weight * (hi0 + hi1)); + } + } + __syncthreads(); + if (threadIdx.x < kRows) + { + float sum = 0.F; +#pragma unroll + for (int k = 0; k < kTopK; ++k) + { + sum += terms[k][threadIdx.x]; + } + out[row0 + threadIdx.x] = __float2bfloat16_rn(sum); + } +#endif +} + +// FC1 + SwiGLU of a multi-token call, one CTA per SM. Consumer warp m < kTokens routes token m as lowMBf16MoeFc1Kernel +// routes its token, the producer warp builds the union of the selections, and the work is the numUnion * interSize +// units of the union: unit u * interSize + j is intermediate index j of union expert u, its gate and up rows. CTA b +// takes units [b * U / G, (b + 1) * U / G) and streams their rows through a ring as lowMBf16MoeFc1Kernel does. x of +// every token is in shared memory, each consumer thread holding there the pieces it multiplies with. Per pair of rows +// a consumer thread loads its pieces of both rows, releases the two buffers, and for each token that routed to the +// rows' expert adds its shares of the two dot products in the order of lowMBf16MoeFc1Kernel. The warp sums of every +// kMultiChunkUnits units go to one of two chunk buffers, from which the epilogue warp forms those units' SwiGLU +// outputs, again in that kernel's order, while the consumer warps fill the other buffer. +template +__global__ void __launch_bounds__(kMultiFc1Threads, 2) lowMBf16MoeMultiFc1Kernel(Bf16* __restrict__ act, + Bf16 const* __restrict__ x, float const* __restrict__ routerLogits, BiasT const* __restrict__ bias, + Bf16 const* __restrict__ w3w1, int hiddenSize, int interSize, int numStages) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + extern __shared__ __align__(128) uint8_t smem[]; + // rowFull[s]: buffer s holds its next row (the producer's arrival plus the copy's bytes). rowFree[s]: every + // consumer warp has read its pieces of buffer s. + __shared__ uint64_t rowFull[kMultiMaxFc1Stages]; + __shared__ uint64_t rowFree[kMultiMaxFc1Stages]; + // chunkFull[b]: every consumer warp has stored its sums of chunk buffer b. chunkFree[b]: the epilogue warp has read + // them. + __shared__ uint64_t chunkFull[2]; + __shared__ uint64_t chunkFree[2]; + __shared__ MultiRouting routing; + + int const warp = threadIdx.x / 32; + int const lane = threadIdx.x % 32; + int const rowBytes = hiddenSize * static_cast(sizeof(Bf16)); + int const piecesPerRow = hiddenSize / kPieceElems; + // ring[numStages][hiddenSize], xShared[kTokens][hiddenSize], rowSums[2][kMultiChunkRows][kTokens][consumer warp] + uint8_t* ring = smem; + uint4* xShared = reinterpret_cast(smem + numStages * rowBytes); + float* rowSums = reinterpret_cast(smem + (numStages + kTokens) * rowBytes); + + if (threadIdx.x == 0) + { + for (int s = 0; s < numStages; ++s) + { + mbarrierInit(&rowFull[s], 1); + mbarrierInit(&rowFree[s], kFc1ConsumerWarps); + } + for (int b = 0; b < 2; ++b) + { + mbarrierInit(&chunkFull[b], kFc1ConsumerWarps); + mbarrierInit(&chunkFree[b], 1); + } + fenceBarrierInit(); + } + for (int e = static_cast(threadIdx.x); e < kNumExperts; e += kMultiFc1Threads) + { + routing.expertTokens[e] = 0; + } + // The correction bias is a constant: the routing warps read it before the grid-dependency wait. + float biasValues[kCandidates] = {}; + if (warp < kTokens) + { + loadBias(bias, lane, biasValues); + } + __syncthreads(); + + // x and the router logits are written by earlier kernels. + cudaGridDependencySynchronize(); + // The FC2 kernel reads only the logits, the bias and its weights before its own wait, and the logits are complete + // since the wait above returned. + cudaTriggerProgrammaticLaunchCompletion(); + + if (warp == kFc1ProducerWarp) + { + namedBarrierSync(kMultiRouteBarrier, kMultiFc1Threads); + buildUnion(routing, lane); + int const numUnion = routing.numUnion; + int const numUnits = numUnion * interSize; + int const unitBegin = static_cast(static_cast(blockIdx.x) * numUnits / gridDim.x); + int const unitEnd = static_cast(static_cast(blockIdx.x + 1) * numUnits / gridDim.x); + int const numRows = 2 * (unitEnd - unitBegin); + + // Lane 0 streams the rows. The next row is of unit (union expert u, index), the gate row first, into buffer s, + // which it fills for the pass-th time. + uint64_t const policy = evictFirstPolicy(); + int64_t const expertStride = 2 * static_cast(interSize) * hiddenSize; + int u = unitBegin / interSize; + int index = unitBegin % interSize; + bool gateRow = true; + Bf16 const* expertRows = w3w1 + routing.unionExperts[u] * expertStride; + int s = 0; + uint32_t pass = 0; + auto issueRow = [&]() + { + int64_t const rowInExpert = gateRow ? interSize + index : index; + mbarrierArriveExpectTx(&rowFull[s], rowBytes); + bulkCopyToShared(ring + s * rowBytes, expertRows + rowInExpert * hiddenSize, rowBytes, &rowFull[s], policy); + if (!gateRow && ++index == interSize) + { + // The up row closed the union expert's last unit. + index = 0; + if (++u < numUnion) + { + expertRows = w3w1 + routing.unionExperts[u] * expertStride; + } + } + gateRow = !gateRow; + if (++s == numStages) + { + s = 0; + ++pass; + } + }; + int const firstRows = numStages < numRows ? numStages : numRows; + if (lane == 0) + { + for (int row = 0; row < firstRows; ++row) + { + issueRow(); + } + } + // Release the consumer and epilogue warps: the union is in shared memory. + __syncwarp(); + namedBarrierArrive(kMultiParkBarrier, kMultiFc1Threads); + if (lane == 0) + { + + for (int row = firstRows; row < numRows; ++row) + { + // Buffer s still holds row - numStages until every consumer warp has read it. + mbarrierWait(&rowFree[s], (pass - 1) & 1U); + issueRow(); + } + } + } + else if (warp == kMultiFc1EpilogueWarp) + { + namedBarrierArrive(kMultiRouteBarrier, kMultiFc1Threads); + namedBarrierSync(kMultiParkBarrier, kMultiFc1Threads); + + int const numUnits = routing.numUnion * interSize; + int const unitBegin = static_cast(static_cast(blockIdx.x) * numUnits / gridDim.x); + int const unitEnd = static_cast(static_cast(blockIdx.x + 1) * numUnits / gridDim.x); + int const numChunks = (unitEnd - unitBegin + kMultiChunkUnits - 1) / kMultiChunkUnits; + for (int chunk = 0; chunk < numChunks; ++chunk) + { + int const buffer = chunk & 1; + mbarrierWait(&chunkFull[buffer], (chunk >> 1) & 1U); + // Unit c of the chunk and token m: the gate row's 8 warp sums, then the up row's, in warp order as + // lowMBf16MoeFc1Kernel adds them. + int const chunkUnitBegin = unitBegin + chunk * kMultiChunkUnits; + int const chunkUnits + = unitEnd - chunkUnitBegin < kMultiChunkUnits ? unitEnd - chunkUnitBegin : kMultiChunkUnits; + float const* bufferSums = rowSums + buffer * kMultiChunkRows * kTokens * kFc1ConsumerWarps; + for (int p = lane; p < chunkUnits * kTokens; p += 32) + { + int const c = p / kTokens; + int const m = p - c * kTokens; + int const unit = chunkUnitBegin + c; + int const unitExpert = unit / interSize; + if ((routing.unionTokens[unitExpert] >> m) & 1U) + { + float const* gateSums = bufferSums + (2 * c * kTokens + m) * kFc1ConsumerWarps; + float const* upSums = bufferSums + ((2 * c + 1) * kTokens + m) * kFc1ConsumerWarps; + float gate = 0.F; + float up = 0.F; + for (int w = 0; w < kFc1ConsumerWarps; ++w) + { + gate += gateSums[w]; + up += upSums[w]; + } + // The CUTLASS FC1 GEMM writes its output in BF16 and the activation kernel reads it back. + gate = roundToBf16(gate); + up = roundToBf16(up); + // silu(gate) * up in FP32 as the CUTLASS activation (GLUAdaptor) computes it, then BF16. + float const silu = gate * (1.F / (1.F + expf(-gate))); + int const slot = routing.unionSlots[unitExpert][m]; + act[(m * kTopK + slot) * interSize + unit - unitExpert * interSize] + = __float2bfloat16_rn(silu * up); + } + } + // The warp's reads of the buffer are done before lane 0 frees it. + __syncwarp(); + if (lane == 0) + { + mbarrierArrive(&chunkFree[buffer]); + } + } + } + else + { + if (warp < kTokens) + { + TopKKey keys[kCandidates]; + buildKeys(routerLogits + warp * kNumExperts, biasValues, lane, keys, nullptr); + int32_t const expert = selectTopK(keys, lane); + if (lane < kTopK) + { + routing.topExperts[warp][lane] = expert; + atomicOr(&routing.expertTokens[expert], 1U << warp); + } + } + // The routing results are in shared memory; the producer warp builds the union from them. + namedBarrierArrive(kMultiRouteBarrier, kMultiFc1Threads); + // Each thread stores its pieces of every token's x; it alone reads them back. + int const piecesPerThread = hiddenSize / kHiddenGranule; + for (int m = 0; m < kTokens; ++m) + { +#pragma unroll + for (int i = 0; i < kMultiMaxPiecesPerThread; ++i) + { + if (i < piecesPerThread) + { + int const piece = m * piecesPerRow + threadIdx.x + i * kFc1ConsumerThreads; + xShared[piece] = __ldcg(reinterpret_cast(x) + piece); + } + } + } + // Wait at the barrier, without polling the ring, while the producer builds the union and issues the first + // rows. + namedBarrierSync(kMultiParkBarrier, kMultiFc1Threads); + int const numUnits = routing.numUnion * interSize; + int const unitBegin = static_cast(static_cast(blockIdx.x) * numUnits / gridDim.x); + int const unitEnd = static_cast(static_cast(blockIdx.x + 1) * numUnits / gridDim.x); + int const numRows = 2 * (unitEnd - unitBegin); + int u = unitBegin / interSize; + int index = unitBegin % interSize; + + int s0 = 0; + uint32_t pass0 = 0; + int chunk = 0; + int chunkRow = 0; + // Rows row and row + 1 are the gate and up rows of unit (u, index), in buffers s0 and s1 (numStages >= 2), + // filled for the pass0-th and pass1-th time. + for (int row = 0; row < numRows; row += 2) + { + int s1 = s0 + 1; + uint32_t pass1 = pass0; + if (s1 == numStages) + { + s1 = 0; + ++pass1; + } + uint32_t const tokens = routing.unionTokens[u]; + mbarrierWait(&rowFull[s0], pass0 & 1U); + mbarrierWait(&rowFull[s1], pass1 & 1U); + uint4 const* pieces0 = reinterpret_cast(ring + s0 * rowBytes); + uint4 const* pieces1 = reinterpret_cast(ring + s1 * rowBytes); + uint4 rowPieces0[kMultiMaxPiecesPerThread]; + uint4 rowPieces1[kMultiMaxPiecesPerThread]; +#pragma unroll + for (int i = 0; i < kMultiMaxPiecesPerThread; ++i) + { + if (i < piecesPerThread) + { + rowPieces0[i] = pieces0[threadIdx.x + i * kFc1ConsumerThreads]; + rowPieces1[i] = pieces1[threadIdx.x + i * kFc1ConsumerThreads]; + } + } + // This warp is done with both buffers (the arrival releases its lanes' reads, ordered before it by the + // warp barrier); the producer refills a buffer once every consumer warp is. + __syncwarp(); + if (lane == 0) + { + mbarrierArrive(&rowFree[s0]); + mbarrierArrive(&rowFree[s1]); + } + if (chunkRow == 0 && chunk >= 2) + { + // The chunk's buffer still holds the sums of chunk - 2 until the epilogue warp has read them. + mbarrierWait(&chunkFree[chunk & 1], ((chunk >> 1) - 1) & 1U); + } + float* chunkSums = rowSums + ((chunk & 1) * kMultiChunkRows + chunkRow) * kTokens * kFc1ConsumerWarps; + for (uint32_t pending = tokens; pending != 0; pending &= pending - 1) + { + int const m = __ffs(static_cast(pending)) - 1; + uint4 const* xPieces = xShared + m * piecesPerRow + threadIdx.x; + uint4 xPiece = xPieces[0]; + float dot0 = dotPiecePacked(rowPieces0[0], xPiece); + float dot1 = dotPiecePacked(rowPieces1[0], xPiece); +#pragma unroll + for (int i = 1; i < kMultiMaxPiecesPerThread; ++i) + { + if (i < piecesPerThread) + { + xPiece = xPieces[i * kFc1ConsumerThreads]; + dot0 += dotPiecePacked(rowPieces0[i], xPiece); + dot1 += dotPiecePacked(rowPieces1[i], xPiece); + } + } + // Lane 0 ends with the gate row's warp sum and lane 16 with the up row's. + float const sum = warpSum2Split(dot0, dot1, lane); + if (lane % 16 == 0) + { + chunkSums[((lane / 16) * kTokens + m) * kFc1ConsumerWarps + warp] = sum; + } + } + s0 = s1 + 1; + pass0 = pass1; + if (s0 == numStages) + { + s0 = 0; + ++pass0; + } + if (++index == interSize) + { + index = 0; + ++u; + } + chunkRow += 2; + if (chunkRow == kMultiChunkRows || row + 2 == numRows) + { + // Lanes 0 and 16 stored this warp's sums of the chunk (ordered before lane 0's arrival by the warp + // barrier); the epilogue warp takes the buffer from here. + __syncwarp(); + if (lane == 0) + { + mbarrierArrive(&chunkFull[chunk & 1]); + } + chunkRow = 0; + ++chunk; + } + } + } +#endif +} + +// FC2 + the weighted sum of a multi-token call. Grid: one CTA per kMultiFc2Rows output rows. Before the grid-dependency +// wait, consumer warp m < kTokens routes token m as lowMBf16MoeFc2Kernel routes its token (same experts, order and +// weights), from the router logits and the bias staged by bulk copies; the producer warp builds the union and starts +// streaming the union experts' weight slices (kMultiFc2Rows rows each) into a ring of numStages stages. After the wait +// the consumer warps stage every token's SwiGLU outputs and take the (union expert, tile) items in turn: the mma of +// lowMBf16MoeFc2Kernel with B column 2 * q + p holding token q's SwiGLU outputs of the expert in the K order of rows of +// parity p, so that accumulator column 2 * q + (g & 1) of row g holds that row's dot product with token q, the same +// products in the same order as lowMBf16MoeFc2Kernel computes for that token alone. Lane (g, t) then holds token t's +// terms of rows g and g + 8 and rounds them as that kernel does; each output row adds its token's 8 terms in slot +// order. +template +__global__ void __launch_bounds__(kMultiFc2Threads, 2) lowMBf16MoeMultiFc2Kernel(Bf16* __restrict__ out, + Bf16 const* __restrict__ act, float const* __restrict__ routerLogits, BiasT const* __restrict__ bias, + Bf16 const* __restrict__ w2, int hiddenSize, int interSize, int numStages, RoutingScale scale) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + constexpr int kGroups = (kTokens + kMultiTokensPerMma - 1) / kMultiTokensPerMma; + constexpr int kProducerWarp = kMultiFc2ConsumerWarps; + extern __shared__ __align__(128) uint8_t smem[]; + __shared__ __align__(16) float logitsShared[kTokens * kNumExperts]; + __shared__ __align__(16) BiasT biasShared[kNumExperts]; + __shared__ uint64_t inputsFull; + // sliceFull[s]: stage s holds its next weight slice. sliceFree[s]: both tile warps have read it. issuedSlices: the + // union experts whose slices the producer warp has issued, in union order. + __shared__ uint64_t sliceFull[kMultiFc2MaxStages]; + __shared__ uint64_t sliceFree[kMultiFc2MaxStages]; + __shared__ int32_t issuedSlices; + __shared__ MultiRouting routing; + __shared__ float topWeights[kTokens][kTopK]; + + int const sliceElems = kMultiFc2Rows * interSize; + // slices[numStages][kMultiFc2Rows][interSize], actShared[kTokens][kTopK][interSize], + // terms[kTokens][kTopK][kMultiFc2Rows] + Bf16* slices = reinterpret_cast(smem); + Bf16* actShared = slices + numStages * sliceElems; + float* terms = reinterpret_cast(actShared + kTokens * kTopK * interSize); + + int const warp = threadIdx.x / 32; + int const lane = threadIdx.x % 32; + int const row0 = blockIdx.x * kMultiFc2Rows; + + if (threadIdx.x == 0) + { + mbarrierInit(&inputsFull, 1); + for (int s = 0; s < numStages; ++s) + { + mbarrierInit(&sliceFull[s], 1); + mbarrierInit(&sliceFree[s], kMultiFc2Tiles); + } + fenceBarrierInit(); + } + for (int e = static_cast(threadIdx.x); e < kNumExperts; e += kMultiFc2Threads) + { + routing.expertTokens[e] = 0; + } + __syncthreads(); + + // Before the grid-dependency wait this kernel reads only the router logits, the bias and the FC2 weights. The + // logits come from the router GEMM, which completed before the FC1 kernel returned from its own wait and released + // this grid; the bias and the weights are constants. The logits and the bias come in by bulk copies, as in + // lowMBf16MoeFc2Kernel. + bool const bulkInputs = ((reinterpret_cast(routerLogits) | reinterpret_cast(bias)) % 16) == 0; + if (warp == kProducerWarp) + { + uint64_t const policy = evictFirstPolicy(); + if (bulkInputs && lane == 0) + { + mbarrierArriveExpectTx(&inputsFull, sizeof(logitsShared) + sizeof(biasShared)); + bulkCopyToShared(logitsShared, routerLogits, sizeof(logitsShared), &inputsFull, policy); + bulkCopyToShared(biasShared, bias, sizeof(biasShared), &inputsFull, policy); + } + namedBarrierSync(kMultiRouteBarrier, kMultiFc2Threads); + buildUnion(routing, lane); + int const numUnion = routing.numUnion; + int const sliceBytes = sliceElems * static_cast(sizeof(Bf16)); + auto issueSlice = [&](int unionIdx, int stage) + { + mbarrierArriveExpectTx(&sliceFull[stage], sliceBytes); + bulkCopyToShared(slices + stage * sliceElems, + w2 + (static_cast(routing.unionExperts[unionIdx]) * hiddenSize + row0) * interSize, sliceBytes, + &sliceFull[stage], policy); + }; + int const firstSlices = numStages < numUnion ? numStages : numUnion; + if (lane == 0) + { + for (int idx = 0; idx < firstSlices; ++idx) + { + issueSlice(idx, idx); + } + issuedSlices = firstSlices; + } + // Release the consumer warps: the union is in shared memory. + __syncwarp(); + namedBarrierArrive(kMultiParkBarrier, kMultiFc2Threads); + // No stage frees up before the consumer warps return from their grid-dependency wait; waiting here instead of + // polling the stages leaves the SM's issue slots to the FC1 CTA beside this one. + cudaGridDependencySynchronize(); + if (lane == 0) + { + for (int idx = numStages; idx < numUnion; ++idx) + { + // Stage idx % numStages still holds union expert idx - numStages until both its tile warps read it. + int const stage = idx % numStages; + mbarrierWait(&sliceFree[stage], ((idx / numStages) - 1) & 1U); + issueSlice(idx, stage); + *reinterpret_cast(&issuedSlices) = idx + 1; + } + } + } + else + { + if (warp < kTokens) + { + float const* tokenLogits = routerLogits + warp * kNumExperts; + if (bulkInputs) + { + mbarrierWait(&inputsFull, 0); + } + TopKKey keys[kCandidates]; +#pragma unroll + for (int i = 0; i < kCandidates; ++i) + { + int const expert = i * 32 + lane; + float const logit + = bulkInputs ? logitsShared[warp * kNumExperts + expert] : __ldcg(tokenLogits + expert); + float const biasValue + = bulkInputs ? static_cast(biasShared[expert]) : static_cast(bias[expert]); + keys[i] = TopKKey{sigmoidAccurate(logit) + biasValue, expert}; + } + int32_t const expert = selectTopK(keys, lane); + __syncwarp(); + // The weights from routingWeight on the selected experts' sigmoid scores, as lowMBf16MoeFc2Kernel forms + // them. + auto const tile = cg::tiled_partition<32>(cg::this_thread_block()); + float const scoreNorm = lane < kTopK + ? sigmoidAccurate(bulkInputs ? logitsShared[warp * kNumExperts + expert] : __ldcg(tokenLogits + expert)) + : 0.F; + float const sumNorm = cg::reduce(tile, scoreNorm, cg::plus{}); + if (lane < kTopK) + { + routing.topExperts[warp][lane] = expert; + topWeights[warp][lane] = routingWeight(scoreNorm, sumNorm, scale); + atomicOr(&routing.expertTokens[expert], 1U << warp); + } + } + // The routing results are in shared memory; the producer warp builds the union from them. + namedBarrierArrive(kMultiRouteBarrier, kMultiFc2Threads); + namedBarrierSync(kMultiParkBarrier, kMultiFc2Threads); + + // The SwiGLU outputs are written by the FC1 kernel. + cudaGridDependencySynchronize(); + int const actPieces = kTokens * kTopK * interSize / kPieceElems; + for (int i = static_cast(threadIdx.x); i < actPieces; i += kMultiFc2ConsumerThreads) + { + reinterpret_cast(actShared)[i] = __ldcg(reinterpret_cast(act) + i); + } + namedBarrierSync(kMultiConsumerBarrier, kMultiFc2ConsumerThreads); + + int const g = lane / 4; + int const t = lane % 4; + int const swap = g & 1; + int const numUnion = routing.numUnion; + for (int item = warp; item < numUnion * kMultiFc2Tiles; item += kMultiFc2ConsumerWarps) + { + int const unionIdx = item / kMultiFc2Tiles; + int const tileIdx = item - unionIdx * kMultiFc2Tiles; + int const stage = unionIdx % numStages; + uint32_t const tokens = routing.unionTokens[unionIdx]; + // B column g of mma group q is token q * kMultiTokensPerMma + g / 2; a token that did not route to this + // expert reads any SwiGLU row, and its accumulator columns are not used. + Bf16 const* actRows[kGroups]; +#pragma unroll + for (int q = 0; q < kGroups; ++q) + { + int const m = q * kMultiTokensPerMma + g / 2; + bool const routed = m < kTokens && ((tokens >> m) & 1U); + actRows[q] = actShared + (routed ? (m * kTopK + routing.unionSlots[unionIdx][m]) * interSize : 0); + } + // The items of a stage go to different warps, so a warp may reach a stage's next fill before the fill + // ahead of it has landed, and a parity wait cannot tell those two fills apart. Once the slice is issued, + // its stage's previous fill has been read, and the parity wait is for this fill. + while (*reinterpret_cast(&issuedSlices) <= unionIdx) + { + __nanosleep(32); + } + mbarrierWait(&sliceFull[stage], (unionIdx / numStages) & 1U); + Bf16 const* slice = slices + stage * sliceElems + tileIdx * kMmaRows * interSize; + float acc[kGroups][2][4] = {}; + for (int j0 = 0; j0 < interSize; j0 += kFc2StepPairK) + { +#pragma unroll + for (int half = 0; half < 2; ++half) + { + int const j = j0 + (half ^ swap) * kStepK + 8 * t; + uint4 const lo = *reinterpret_cast(slice + g * interSize + j); + uint4 const hi = *reinterpret_cast(slice + (g + 8) * interSize + j); +#pragma unroll + for (int q = 0; q < kGroups; ++q) + { + uint4 const av = *reinterpret_cast(actRows[q] + j); + mmaStep(acc[q][half], lo, hi, av); + } + } + } + // This warp is done with the stage; the producer refills it once both tile warps are. + __syncwarp(); + if (lane == 0) + { + mbarrierArrive(&sliceFree[stage]); + } +#pragma unroll + for (int q = 0; q < kGroups; ++q) + { + int const m = q * kMultiTokensPerMma + t; + if (m < kTokens && ((tokens >> m) & 1U)) + { + int const slot = routing.unionSlots[unionIdx][m]; + float const weight = topWeights[m][slot]; + // Column 2 * t + (g & 1) of rows g and g + 8, as lowMBf16MoeFc2Kernel's lane (g, 0) takes column + // g & 1. + float const lo0 = swap ? acc[q][0][1] : acc[q][0][0]; + float const lo1 = swap ? acc[q][1][1] : acc[q][1][0]; + float const hi0 = swap ? acc[q][0][3] : acc[q][0][2]; + float const hi1 = swap ? acc[q][1][3] : acc[q][1][2]; + float* tokenTerms = terms + (m * kTopK + slot) * kMultiFc2Rows + tileIdx * kMmaRows; + tokenTerms[g] = roundToBf16(weight * (lo0 + lo1)); + tokenTerms[g + 8] = roundToBf16(weight * (hi0 + hi1)); + } + } + } + // Every term is in shared memory: each output row adds its token's 8 terms in slot order. + namedBarrierSync(kMultiConsumerBarrier, kMultiFc2ConsumerThreads); + for (int i = static_cast(threadIdx.x); i < kTokens * kMultiFc2Rows; i += kMultiFc2ConsumerThreads) + { + int const m = i / kMultiFc2Rows; + int const r = i - m * kMultiFc2Rows; + float sum = 0.F; +#pragma unroll + for (int k = 0; k < kTopK; ++k) + { + sum += terms[(m * kTopK + k) * kMultiFc2Rows + r]; + } + out[m * hiddenSize + row0 + r] = __float2bfloat16_rn(sum); + } + } +#endif +} + +// Lets a kernel use smemBytes of dynamic shared memory and asks for the largest shared-memory carveout for both +// kernels, so that an SM running a CTA of one can take a CTA of the other. +template +void setSharedMemoryAttributes(Kernel* kernel, size_t smemBytes) +{ + TLLM_CUDA_CHECK( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smemBytes))); + TLLM_CUDA_CHECK(cudaFuncSetAttribute( + kernel, cudaFuncAttributePreferredSharedMemoryCarveout, static_cast(cudaSharedmemCarveoutMaxShared))); +} + +// routingWeight's FP32 path applies when the scaling factor is a float with at most 20 significant bits (the low 4 +// mantissa bits zero) between 2^-20 and 2^20. +RoutingScale makeRoutingScale(double routedScalingFactor) +{ + float const valueFloat = static_cast(routedScalingFactor); + uint32_t bits = 0; + std::memcpy(&bits, &valueFloat, sizeof(bits)); + bool const fastPath = static_cast(valueFloat) == routedScalingFactor && (bits & 0xFU) == 0 + && valueFloat >= 0x1p-20f && valueFloat <= 0x1p20f; + return RoutingScale{routedScalingFactor, valueFloat, fastPath ? 1 : 0}; +} + +template +void launchFc2(Bf16* out, Bf16 const* act, float const* routerLogits, BiasT const* bias, Bf16 const* w2, int hiddenSize, + int interSize, double routedScalingFactor, cudaStream_t stream) +{ + int const rows = kTiles * kMmaRows; + size_t const smemBytes = static_cast(kTopK) * (rows + 1) * interSize * sizeof(Bf16); + auto* kernel = &lowMBf16MoeFc2Kernel; + setSharedMemoryAttributes(kernel, smemBytes); + common::launchWithPdlWhenEnabled("lowMBf16MoeFc2", kernel, dim3(hiddenSize / rows), dim3(kFc2Threads), smemBytes, + stream, out, act, routerLogits, bias, w2, hiddenSize, interSize, makeRoutingScale(routedScalingFactor)); +} + +// The shared memory of an SM, the per-CTA reservation and the static shared memory of the two multi-token kernels. +struct MultiSmemLimits +{ + size_t perSm; + size_t reservedPerCta; + size_t fc1Static; + size_t fc2Static; +}; + +template +MultiSmemLimits queryMultiSmemLimits(Fc1Kernel* fc1, Fc2Kernel* fc2) +{ + int device = 0; + TLLM_CUDA_CHECK(cudaGetDevice(&device)); + int perSm = 0; + int reserved = 0; + TLLM_CUDA_CHECK(cudaDeviceGetAttribute(&perSm, cudaDevAttrMaxSharedMemoryPerMultiprocessor, device)); + TLLM_CUDA_CHECK(cudaDeviceGetAttribute(&reserved, cudaDevAttrReservedSharedMemoryPerBlock, device)); + cudaFuncAttributes fc1Attributes{}; + cudaFuncAttributes fc2Attributes{}; + TLLM_CUDA_CHECK(cudaFuncGetAttributes(&fc1Attributes, fc1)); + TLLM_CUDA_CHECK(cudaFuncGetAttributes(&fc2Attributes, fc2)); + return MultiSmemLimits{static_cast(perSm), static_cast(reserved), fc1Attributes.sharedSizeBytes, + fc2Attributes.sharedSizeBytes}; +} + +template +void launchLowMBf16MoeMulti(Bf16* out, Bf16* act, Bf16 const* x, float const* routerLogits, BiasT const* bias, + Bf16 const* w3w1, Bf16 const* w2, int hiddenSize, int interSize, double routedScalingFactor, cudaStream_t stream) +{ + auto* fc1 = &lowMBf16MoeMultiFc1Kernel; + auto* fc2 = &lowMBf16MoeMultiFc2Kernel; + static int const smCount = common::getMultiProcessorCount(); + static MultiSmemLimits const limits = queryMultiSmemLimits(fc1, fc2); + + // Shared memory per CTA as an SM accounts it: the FC1 CTA without its ring, and the FC2 CTA without its stages. + int const rowBytes = hiddenSize * static_cast(sizeof(Bf16)); + size_t const sliceBytes = static_cast(kMultiFc2Rows) * interSize * sizeof(Bf16); + size_t const fc1Fixed = static_cast(kTokens) * rowBytes + + static_cast(2 * kMultiChunkRows * kTokens * kFc1ConsumerWarps) * sizeof(float); + size_t const fc2Fixed = static_cast(kTokens * kTopK) * interSize * sizeof(Bf16) + + static_cast(kTokens * kTopK * kMultiFc2Rows) * sizeof(float); + size_t const fc1Base = limits.fc1Static + limits.reservedPerCta + fc1Fixed; + size_t const fc2Base = limits.fc2Static + limits.reservedPerCta + fc2Fixed; + + // FC1: one CTA per SM (each with at least one unit: the union has at least kTopK experts). Its ring takes enough + // rows that two FC1 CTAs never share an SM (the static unit split would give that SM twice the rows), and at most + // as many as leave room for an FC2 CTA with two stages beside it. + size_t const halfSm = limits.perSm / 2 + 1; + int const oneCtaStages + = halfSm > fc1Base ? static_cast((halfSm - fc1Base + static_cast(rowBytes) - 1) / rowBytes) : 0; + size_t const fc2Min = fc2Base + 2 * sliceBytes; + int const roomStages + = limits.perSm > fc1Base + fc2Min ? static_cast((limits.perSm - fc1Base - fc2Min) / rowBytes) : 0; + int numStages = std::min(kMultiMaxFc1Stages, std::max(kMultiMinFc1Stages, oneCtaStages)); + if (roomStages >= 2) + { + numStages = std::min(numStages, roomStages); + } + int const fc1Grid = std::min(smCount, kTopK * interSize); + size_t const fc1SmemBytes = static_cast(numStages) * rowBytes + fc1Fixed; + setSharedMemoryAttributes(fc1, fc1SmemBytes); + common::launchWithPdlWhenEnabled("lowMBf16MoeMultiFc1", fc1, dim3(fc1Grid), dim3(kMultiFc1Threads), fc1SmemBytes, + stream, act, x, routerLogits, bias, w3w1, hiddenSize, interSize, numStages); + + // FC2: as many weight stages as fit beside the FC1 CTA and leave room for two FC2 CTAs per SM once the FC1 CTAs + // exit, at least 2. + size_t const fc1Total = fc1Base + static_cast(numStages) * rowBytes; + size_t const fc2Room = std::min(limits.perSm - std::min(limits.perSm, fc1Total), limits.perSm / 2); + int const fitStages = fc2Room > fc2Base ? static_cast((fc2Room - fc2Base) / sliceBytes) : 0; + int const fc2Stages = std::max(2, std::min(kMultiFc2MaxStages, fitStages)); + size_t const fc2SmemBytes = static_cast(fc2Stages) * sliceBytes + fc2Fixed; + setSharedMemoryAttributes(fc2, fc2SmemBytes); + common::launchWithPdlWhenEnabled("lowMBf16MoeMultiFc2", fc2, dim3(hiddenSize / kMultiFc2Rows), + dim3(kMultiFc2Threads), fc2SmemBytes, stream, out, act, routerLogits, bias, w2, hiddenSize, interSize, + fc2Stages, makeRoutingScale(routedScalingFactor)); +} + +} // namespace + +bool lowMBf16MoeSupportsShape(int64_t hiddenSize, int64_t interSize) +{ + return hiddenSize > 0 && hiddenSize % kHiddenGranule == 0 && hiddenSize <= kMaxPiecesPerThread * kHiddenGranule + && interSize > 0 && interSize % kFc2StepPairK == 0 && interSize <= kMaxInterSize; +} + +template +void invokeLowMBf16Moe(Bf16* out, Bf16* act, Bf16 const* x, float const* routerLogits, BiasT const* bias, + Bf16 const* w3w1, Bf16 const* w2, int hiddenSize, int interSize, double routedScalingFactor, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(lowMBf16MoeSupportsShape(hiddenSize, interSize), + "lowMBf16Moe: unsupported hidden size %d or intermediate size %d", hiddenSize, interSize); + + // FC1: one CTA per SM, each with an even share of the units and a slot per row and consumer warp for the sums. + static int const smCount = common::getMultiProcessorCount(); + int const numUnits = kTopK * interSize; + int const fc1Grid = std::min(smCount, numUnits); + int const rowBytes = hiddenSize * static_cast(sizeof(Bf16)); + int const numStages = std::min(kMaxStages, kFc1RingBytes / rowBytes); + int const maxRows = 2 * ((numUnits + fc1Grid - 1) / fc1Grid); + size_t const fc1SmemBytes + = static_cast(numStages) * rowBytes + static_cast(maxRows) * kFc1ConsumerWarps * sizeof(float); + auto* fc1 = &lowMBf16MoeFc1Kernel; + setSharedMemoryAttributes(fc1, fc1SmemBytes); + common::launchWithPdlWhenEnabled("lowMBf16MoeFc1", fc1, dim3(fc1Grid), dim3(kFc1Threads), fc1SmemBytes, stream, act, + x, routerLogits, bias, w3w1, hiddenSize, interSize, numStages); + + // Two 16-row tiles per FC2 CTA where the 8 weight slices fit in shared memory, else one. + if (interSize <= kMaxInterSize / 2) + { + launchFc2(out, act, routerLogits, bias, w2, hiddenSize, interSize, routedScalingFactor, stream); + } + else + { + launchFc2(out, act, routerLogits, bias, w2, hiddenSize, interSize, routedScalingFactor, stream); + } +} + +bool lowMBf16MoeMultiSupportsShape(int64_t numTokens, int64_t hiddenSize, int64_t interSize) +{ + return numTokens >= 2 && numTokens <= kMultiMaxTokens && hiddenSize > 0 && hiddenSize % kHiddenGranule == 0 + && hiddenSize <= kMultiMaxPiecesPerThread * kHiddenGranule && interSize > 0 && interSize % kFc2StepPairK == 0 + && interSize <= kMaxInterSize; +} + +template +void invokeLowMBf16MoeMulti(Bf16* out, Bf16* act, Bf16 const* x, float const* routerLogits, BiasT const* bias, + Bf16 const* w3w1, Bf16 const* w2, int numTokens, int hiddenSize, int interSize, double routedScalingFactor, + cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(lowMBf16MoeMultiSupportsShape(numTokens, hiddenSize, interSize), + "lowMBf16MoeMulti: unsupported token count %d, hidden size %d or intermediate size %d", numTokens, hiddenSize, + interSize); + switch (numTokens) + { + case 2: + launchLowMBf16MoeMulti( + out, act, x, routerLogits, bias, w3w1, w2, hiddenSize, interSize, routedScalingFactor, stream); + break; + case 3: + launchLowMBf16MoeMulti( + out, act, x, routerLogits, bias, w3w1, w2, hiddenSize, interSize, routedScalingFactor, stream); + break; + case 4: + launchLowMBf16MoeMulti( + out, act, x, routerLogits, bias, w3w1, w2, hiddenSize, interSize, routedScalingFactor, stream); + break; + case 5: + launchLowMBf16MoeMulti( + out, act, x, routerLogits, bias, w3w1, w2, hiddenSize, interSize, routedScalingFactor, stream); + break; + default: + launchLowMBf16MoeMulti( + out, act, x, routerLogits, bias, w3w1, w2, hiddenSize, interSize, routedScalingFactor, stream); + break; + } +} + +#define INSTANTIATE_LOW_M_BF16_MOE(BiasT) \ + template void invokeLowMBf16Moe(Bf16 * out, Bf16 * act, Bf16 const* x, float const* routerLogits, \ + BiasT const* bias, Bf16 const* w3w1, Bf16 const* w2, int hiddenSize, int interSize, \ + double routedScalingFactor, cudaStream_t stream); + +template void invokeLowMBf16MoeMulti(Bf16*, Bf16*, Bf16 const*, float const*, float const*, Bf16 const*, + Bf16 const*, int, int, int, double, cudaStream_t); +template void invokeLowMBf16MoeMulti<__nv_bfloat16>(Bf16*, Bf16*, Bf16 const*, float const*, __nv_bfloat16 const*, + Bf16 const*, Bf16 const*, int, int, int, double, cudaStream_t); + +INSTANTIATE_LOW_M_BF16_MOE(float); +INSTANTIATE_LOW_M_BF16_MOE(__nv_bfloat16); + +#undef INSTANTIATE_LOW_M_BF16_MOE + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.h b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmLowMBf16Moe.h similarity index 61% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.h rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmLowMBf16Moe.h index 928eebda65e2..92098b5fadd5 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.h +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmLowMBf16Moe.h @@ -69,6 +69,40 @@ void invokeLowMBf16Moe(__nv_bfloat16* out, __nv_bfloat16* act, __nv_bfloat16 con BiasT const* bias, __nv_bfloat16 const* w3w1, __nv_bfloat16 const* w2, int hiddenSize, int interSize, double routedScalingFactor, cudaStream_t stream); +//! Most tokens invokeLowMBf16MoeMulti takes. +constexpr int kLowMBf16MoeMaxTokens = 6; + +//! Whether invokeLowMBf16MoeMulti supports a layer and token count: 2 to kLowMBf16MoeMaxTokens tokens, hiddenSize a +//! multiple of 2048 and at most 6144, interSize a multiple of 64 and at most 512. +bool lowMBf16MoeMultiSupportsShape(int64_t numTokens, int64_t hiddenSize, int64_t interSize); + +//! invokeLowMBf16Moe for 2 to kLowMBf16MoeMaxTokens tokens, in two kernels: +//! 1. route every token, then stream the FC1 rows of the union of the routed experts once, and for each token that +//! routed to the expert: act[m, k, :] = bf16(silu(bf16(W1[e] x_m)) * bf16(W3[e] x_m)), k the expert's slot in +//! token m's list; +//! 2. route again, stream the FC2 rows of the union, and out[m] = bf16(sum_k bf16(w_mk * (W2[e_mk] act[m, k, :]))). +//! Each output row is bit-identical to invokeLowMBf16Moe on that token alone: the same routing, the same per-row +//! reductions in the same order and the same rounding points. The union is in ascending expert order in both kernels. +//! +//! The FC1 kernel runs one CTA per SM, as invokeLowMBf16Moe's does, with x of every token in shared memory; its units +//! are (union expert, intermediate index) pairs, and an epilogue warp forms the SwiGLU outputs of each chunk of units +//! while the consumer warps go on. The FC2 kernel takes 16 output rows per CTA: one CTA fits beside each FC1 CTA, +//! routes and stages the FC2 weights of the first union experts in a ring of shared-memory stages before its +//! grid-dependency wait, and two share an SM once the FC1 CTAs exit; the ring streams the rest of the union after the +//! wait. +//! +//! Under PDL the FC1 kernel reads x and the logits after its grid-dependency wait and releases the FC2 kernel right +//! after it; the FC2 kernel reads only the logits, the bias and the FC2 weights before its own wait. +//! +//! \param out Output [numTokens, hiddenSize]. +//! \param act Scratch [numTokens, kLowMBf16MoeTopK, interSize] for the SwiGLU outputs, 16-byte aligned. +//! \param x Input tokens [numTokens, hiddenSize], 16-byte aligned. +//! \param routerLogits Router logits [numTokens, kLowMBf16MoeNumExperts], FP32. +template +void invokeLowMBf16MoeMulti(__nv_bfloat16* out, __nv_bfloat16* act, __nv_bfloat16 const* x, float const* routerLogits, + BiasT const* bias, __nv_bfloat16 const* w3w1, __nv_bfloat16 const* w2, int numTokens, int hiddenSize, int interSize, + double routedScalingFactor, cudaStream_t stream); + } // namespace kernels::glm_moe_dsa TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.cu b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMlaKvHalfAppend.cu similarity index 99% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.cu rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMlaKvHalfAppend.cu index 91246637cf1f..790461f94431 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.cu +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMlaKvHalfAppend.cu @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmMlaKvHalfAppend.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.h b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMlaKvHalfAppend.h similarity index 100% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.h rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMlaKvHalfAppend.h diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.cu b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMnnvlAllgather.cu similarity index 98% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.cu rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMnnvlAllgather.cu index c2753b020af7..45e379ce4436 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.cu +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMnnvlAllgather.cu @@ -15,13 +15,13 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmMnnvlAllgather.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/common/lamportUtils.cuh" -#include "tensorrt_llm/kernels/glmMoeDsa/glmLamport.cuh" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmLamport.cuh" #include #include diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.h b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMnnvlAllgather.h similarity index 100% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.h rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMnnvlAllgather.h diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.cu b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMnnvlHop.cu similarity index 99% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.cu rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMnnvlHop.cu index 6a1ca204bf1a..dadc93b4b9d1 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.cu +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMnnvlHop.cu @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmMnnvlHop.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaTypeUtils.cuh" @@ -23,7 +23,7 @@ #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/common/lamportUtils.cuh" #include "tensorrt_llm/common/reduceKernelUtils.cuh" -#include "tensorrt_llm/kernels/glmMoeDsa/glmLamport.cuh" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmLamport.cuh" #include #include diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.h b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMnnvlHop.h similarity index 100% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.h rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMnnvlHop.h diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeFusedRouter.cu b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeFusedRouter.cu new file mode 100644 index 000000000000..aea82418b195 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeFusedRouter.cu @@ -0,0 +1,1415 @@ +/* + * 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. + */ + +// Router GEMV + DeepSeek-style routing (sigmoid + bias, top-8 of 256 experts, scaled sum normalization) in one kernel +// for 1-8 tokens, with the outputs of routingIndicesSmallTokenKernel. +// +// Logits. Every logit is the value glmRouterGemm.cu's router_gemm_kernel computes (bit for bit). There, thread t +// (0..127) of the CTA of expert e accumulates acc = fma(a, w, acc) over k = ki * 1024 + 8 t + j (ki = 0..5, j = 0..7), +// in that order; warp s (the "slice" of threads 32 s .. 32 s + 31) adds its lanes with the xor butterfly 16, 8, 4, 2, 1 +// (lane 0's value), and the four slice sums are added to 0.0f in slice order. Here a GEMV CTA covers one slice of 8 +// experts: it stages the slice's activations of every token in shared memory (24 KB at 8 tokens), runs the same +// per-thread chains (fma.rn.f32.bf16 computes fma(float(a), float(w), acc) without conversion instructions), adds the +// 32 positions with the same pairs as the butterfly and publishes the slice sums. 128 CTAs cover 256 experts x 4 +// slices, and each CTA stages a quarter of the activations instead of all of them. +// +// Handoff. The slice sums go to one of two exchange buffers in the workspace (by the parity of the epoch word) whose +// idle value is -0.0. A sum of -0.0 is published as +0.0, which gives the same logit, since the four sums are added to +// 0.0f. The routing CTA polls its token rows until no -0.0 is left, so no fence or counter is needed. It resets the +// other buffer, consumed by the previous launch, while it waits, and advances the epoch at its end. Launches that share +// a workspace must therefore run in stream order. +// +// Routing. Phase A of routingIndicesSmallTokenKernel (one warp per token: sigmoid + bias, the packed-key top-K, the +// scaled sum normalization; the same operations in the same order). The metadata then follows from the set of +// selected local experts after a single barrier: with no more tokens than tile rows every selected local expert owns +// exactly one CTA tile, so its tile index is its rank among the selected local experts and its rows start at +// rank * tile, which is what the small-token kernel's scans compute. + +#include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeFusedRouter.h" +#include "tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingCustomPolicy.cuh" +#include "tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingKernel.h" + +#include + +// The kernel is file-local to the routing namespace whose policies and helpers it uses. +namespace moe::dev::routing +{ +namespace routingCustom +{ +namespace +{ + +using FusedExpertSelect = TopKExpertSelect; +using FusedKernelParams = KernelParams; + +constexpr int kNumExperts = 256; +constexpr int kTopK = 8; +constexpr int kMaxTokens = 8; +constexpr int kHiddenDim = 6144; +constexpr int kEltsPerVec = 8; // bf16 per 16-byte load, as in the router GEMV +constexpr int kKIterElts = kEltsPerVec * 128; // elements per k-iteration of the router GEMV +constexpr int kKIters = kHiddenDim / kKIterElts; // 6 +constexpr int kSlices = 4; // warps per expert in the router GEMV +constexpr int kExpertsPerCta = 8; +constexpr int kNumGemvCtas = kNumExperts / kExpertsPerCta * kSlices; // 128 +constexpr int kNumThreads = 256; +constexpr int kSliceVecs = kKIters * WarpSize; // 16-byte vectors of one token's slice +constexpr int kActBytes = kMaxTokens * kSliceVecs * 16; // 24 KB +constexpr int kPartBytes = WarpSize * kExpertsPerCta * kMaxTokens * 4; // 8 KB +constexpr int kSmemBytes = kActBytes + kPartBytes; +constexpr int kExchangeFloats = kMaxTokens * kNumExperts * kSlices; +constexpr int kKIdxRowBytes = 8; +constexpr uint32_t kNegZeroBits = 0x80000000u; + +static_assert(kNumExperts == kNumThreads, "the routing CTA runs one thread per expert"); +static_assert( + 2 * kExchangeFloats + 4 == tensorrt_llm::kernels::glm_moe_dsa::kFusedRouterWorkspaceWords, "workspace layout"); + +struct FusedRouterParams +{ + FusedKernelParams routing; + __nv_bfloat16 const* hiddenStates; + __nv_bfloat16 const* routerWeight; + int32_t* workspace; +}; + +__device__ __forceinline__ uint32_t smemAddr(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ void gridDependencyWait() +{ + // The memory clobber keeps the compiler from moving memory accesses across the wait. + asm volatile("griddepcontrol.wait;" ::: "memory"); +} + +// Non-coherent 16-byte load of router weights. +__device__ __forceinline__ uint4 loadWeight16(void const* ptr) +{ + uint4 v; + asm volatile("ld.global.nc.v4.u32 {%0, %1, %2, %3}, [%4];" : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(ptr)); + return v; +} + +__device__ __forceinline__ uint32_t loadVolatile(uint32_t const* ptr) +{ + uint32_t v; + asm volatile("ld.volatile.global.u32 %0, [%1];" : "=r"(v) : "l"(ptr) : "memory"); + return v; +} + +__device__ __forceinline__ float4 loadVolatile4(float const* ptr) +{ + float4 v; + asm volatile("ld.volatile.global.v4.f32 {%0, %1, %2, %3}, [%4];" + : "=f"(v.x), "=f"(v.y), "=f"(v.z), "=f"(v.w) + : "l"(ptr) + : "memory"); + return v; +} + +// A slice sum for the routing CTA, which polls it with volatile loads. +__device__ __forceinline__ void storeRelaxed(float* ptr, float v) +{ + asm volatile("st.relaxed.gpu.global.f32 [%0], %1;" ::"l"(ptr), "f"(v) : "memory"); +} + +__device__ __forceinline__ bool isNegZero(float v) +{ + return __float_as_uint(v) == kNegZeroBits; +} + +// acc = fma(a[j], w[j], acc) for j = 0..7, in that order, on 8 packed bf16 pairs. +__device__ __forceinline__ void fmaBf16x8(float& acc, uint4 const& a, uint4 const& w) +{ + uint32_t const aw[4] = {a.x, a.y, a.z, a.w}; + uint32_t const ww[4] = {w.x, w.y, w.z, w.w}; +#pragma unroll + for (int q = 0; q < 4; q++) + { + asm("{\n" + ".reg .b16 al, ah, bl, bh;\n" + "mov.b32 {al, ah}, %1;\n" + "mov.b32 {bl, bh}, %2;\n" + "fma.rn.f32.bf16 %0, al, bl, %0;\n" + "fma.rn.f32.bf16 %0, ah, bh, %0;\n" + "}\n" + : "+f"(acc) + : "r"(aw[q]), "r"(ww[q])); + } +} + +// The slice GEMV of one CTA: slice r of experts 8 g .. 8 g + 7. Thread mapping: warp w covers positions 8 (w & 3) .. +// + 7 of the slice and experts 4 (w >> 2) .. + 3, so the 4 expert groups of a warp read the same 128 bytes of +// activations. Returns through emit(eLocal, m, sliceSum) once per (expert, token). +template +__device__ __forceinline__ void sliceGemv(uint8_t* smem, __nv_bfloat16 const* hiddenStates, uint4 const (&w)[kKIters], + int r, int t, int eLocal, Emit const& emit) +{ + uint4* act = reinterpret_cast(smem); + float* part = reinterpret_cast(smem + kActBytes); + + // Stage the slice: token m, k-iteration ki, positions 32 r .. 32 r + 31. + constexpr int kVecs = kNumTokens * kSliceVecs; +#pragma unroll + for (int j = 0; j < (kVecs + kNumThreads - 1) / kNumThreads; ++j) + { + int const c = threadIdx.x + kNumThreads * j; + if (c < kVecs) + { + int const m = c / kSliceVecs; + int const rem = c - m * kSliceVecs; + __nv_bfloat16 const* src = hiddenStates + m * kHiddenDim + (rem / WarpSize) * kKIterElts + + (WarpSize * r + rem % WarpSize) * kEltsPerVec; + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ::"r"(smemAddr(&act[c])), "l"(src) : "memory"); + } + } + asm volatile("cp.async.commit_group;\ncp.async.wait_group 0;" ::: "memory"); + __syncthreads(); + + // The per-thread chains of router thread 32 r + t for one expert and every token. + float acc[kMaxTokens] = {}; +#pragma unroll + for (int ki = 0; ki < kKIters; ki++) + { +#pragma unroll + for (int m = 0; m < kNumTokens; m++) + { + fmaBf16x8(acc[m], act[(m * kKIters + ki) * WarpSize + t], w[ki]); + } + } + + // The butterfly of the slice: 4 threads per (expert, token) fold the positions {sub, sub + 4, ..., sub + 28} + // through levels 16, 8 and 4, two shuffles add levels 2 and 1. + float4* row = reinterpret_cast(part + (t * kExpertsPerCta + eLocal) * kMaxTokens); + row[0] = make_float4(acc[0], acc[1], acc[2], acc[3]); + row[1] = make_float4(acc[4], acc[5], acc[6], acc[7]); + __syncthreads(); + int const pair = threadIdx.x / 4; + int const sub = threadIdx.x % 4; + int const eL = pair / kMaxTokens; + int const m = pair % kMaxTokens; + auto v = [&](int tt) { return part[(tt * kExpertsPerCta + eL) * kMaxTokens + m]; }; + float const s1a = v(sub) + v(sub + 16); + float const s1b = v(sub + 4) + v(sub + 20); + float const s1c = v(sub + 8) + v(sub + 24); + float const s1d = v(sub + 12) + v(sub + 28); + float sum = (s1a + s1c) + (s1b + s1d); + sum += __shfl_xor_sync(0xffffffff, sum, 2); + sum += __shfl_xor_sync(0xffffffff, sum, 1); + if (sub == 0 && m < kNumTokens) + { + emit(eL, m, sum); + } +} + +template +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) +__global__ void __launch_bounds__(kNumThreads, 1) routingFusedRouterKernel(FusedRouterParams const __grid_constant__ p) +{ + extern __shared__ __align__(128) uint8_t smemSlice[]; + __shared__ __align__(8) uint8_t smemKIdx[kNumExperts * kKIdxRowBytes]; + __shared__ int16_t smemTopKExpert[kMaxTokens * kTopK]; + __shared__ float smemBias[kNumExperts]; + __shared__ __align__(16) uint32_t smemSelected[kNumExperts / 32]; + + auto const& params = p.routing; + float* const buffers = reinterpret_cast(p.workspace); + uint32_t* const epochPtr = reinterpret_cast(p.workspace + 2 * kExchangeFloats); + + if (blockIdx.x < kNumGemvCtas) + { + int const r = blockIdx.x % kSlices; + int const g = blockIdx.x / kSlices; + int const warp = threadIdx.x / WarpSize; + int const lane = threadIdx.x % WarpSize; + int const t = 8 * (warp % 4) + lane % 8; + int const eLocal = 4 * (warp / 4) + lane / 8; + // The weights are not written by the primary grid: load them before the wait. A shared-memory store of their + // XOR ahead of a barrier makes the loads complete before the wait instead of being scheduled after it. + uint4 w[kKIters]; + __nv_bfloat16 const* wRow + = p.routerWeight + (kExpertsPerCta * g + eLocal) * kHiddenDim + (WarpSize * r + t) * kEltsPerVec; + uint32_t loaded = 0; +#pragma unroll + for (int ki = 0; ki < kKIters; ki++) + { + w[ki] = loadWeight16(wRow + ki * kKIterElts); + loaded ^= w[ki].x ^ w[ki].y ^ w[ki].z ^ w[ki].w; + } + reinterpret_cast(smemSlice + kActBytes)[threadIdx.x] = loaded; + __syncthreads(); + gridDependencyWait(); + // The previous launch on this workspace precedes the primary grid in the stream, so it has completed. + float* const exchange = buffers + (loadVolatile(epochPtr) & 1u) * kExchangeFloats; + sliceGemv(smemSlice, p.hiddenStates, w, r, t, eLocal, + [&](int eL, int m, float sliceSum) + { + storeRelaxed(&exchange[(m * kNumExperts + kExpertsPerCta * g + eL) * kSlices + r], + isNegZero(sliceSum) ? 0.f : sliceSum); + }); + // Exit without a trigger: the dependents launch once the routing CTA triggers. + return; + } + + // Routing CTA. Independent of the primary grid: clear this expert's row of the byte map, stage its bias. + *reinterpret_cast(&smemKIdx[threadIdx.x * kKIdxRowBytes]) = uint2{0xFFFFFFFFu, 0xFFFFFFFFu}; + auto const& pre = params.mExpertSelectParams.mPreprocessParams; + smemBias[threadIdx.x] = loadScalar(pre.ptrRoutingBias, threadIdx.x, pre.dtypeBias); + if (threadIdx.x < kNumExperts / 32) + { + smemSelected[threadIdx.x] = 0u; + } + __syncthreads(); + gridDependencyWait(); + uint32_t const epoch = loadVolatile(epochPtr); + float* const exchange = buffers + (epoch & 1u) * kExchangeFloats; + { + // The buffer of the previous launch, which its routing CTA has consumed: idle again for the next launch. + float4* other = reinterpret_cast(buffers + ((epoch + 1u) & 1u) * kExchangeFloats); + for (int i = threadIdx.x; i < kExchangeFloats / 4; i += kNumThreads) + { + other[i] = make_float4(-0.f, -0.f, -0.f, -0.f); + } + } + + // ----- Phase A: one warp per token, as in routingIndicesSmallTokenKernel ----- + int32_t const warpIdx = __shfl_sync(0xffffffff, threadIdx.x / WarpSize, 0); + int32_t const laneIdx = cutlass::arch::LaneId(); + auto block = cg::this_thread_block(); + auto warp = cg::tiled_partition(block); + if (warpIdx < kNumTokens) + { + constexpr int VecSize = kNumExperts / WarpSize; + float4 slices[VecSize]; + bool ready; + do + { + ready = true; +#pragma unroll + for (int i = 0; i < VecSize; ++i) + { + slices[i] = loadVolatile4(exchange + (warpIdx * kNumExperts + i * WarpSize + laneIdx) * kSlices); + ready = ready && !isNegZero(slices[i].x) && !isNegZero(slices[i].y) && !isNegZero(slices[i].z) + && !isNegZero(slices[i].w); + } + } while (!ready); + + float score[VecSize]; + int32_t idx[VecSize]; +#pragma unroll + for (int i = 0; i < VecSize; ++i) + { + idx[i] = i * WarpSize + laneIdx; + float logit = 0.0f; + logit += slices[i].x; + logit += slices[i].y; + logit += slices[i].z; + logit += slices[i].w; + score[i] = sigmoid_accurate(logit) + smemBias[idx[i]]; + } + float topScore[kTopK]; + int32_t topExpert[kTopK]; + topk::reduceTopK(warp, topScore, topExpert, score, idx, float{-INFINITY}, kTopK); + + // Lane k's top-K entry, the value the small-token kernel reads as warpTopKScore[laneIdx]. + float myScore = 0.f; + int32_t myExpert = 0; +#pragma unroll + for (int k = 0; k < kTopK; ++k) + { + myScore = laneIdx == k ? topScore[k] : myScore; + myExpert = laneIdx == k ? topExpert[k] : myExpert; + } + bool const isTop = laneIdx < kTopK; + float const sigmoidScore = isTop ? (myScore - smemBias[myExpert]) : 0.f; + // cg::reduce's butterfly (xor 16, 8, 4, 2, 1) adds +0.0 to lanes 0-7 in its first two levels, which leaves + // them unchanged (no sigmoidScore is -0.0); levels 4, 2 and 1 then pair lanes 0-7 as here. + float sum = sigmoidScore; + sum += __shfl_xor_sync(0xffffffff, sum, 4); + sum += __shfl_xor_sync(0xffffffff, sum, 2); + sum += __shfl_xor_sync(0xffffffff, sum, 1); + if (isTop) + { + auto const& post = params.mExpertSelectParams.mPostprocessParams; + float const weight = sigmoidScore * post.routeScale / (sum + post.sumEpsilon); + smemKIdx[myExpert * kKIdxRowBytes + warpIdx] = static_cast(laneIdx); + smemTopKExpert[warpIdx * kTopK + laneIdx] = static_cast(myExpert); + atomicOr(&smemSelected[myExpert / 32], 1u << (myExpert % 32)); + params.mPtrTopKWeights[warpIdx * kTopK + laneIdx] = __nv_bfloat16{weight}; + } + } + __syncthreads(); + + // ----- Metadata: tile index = rank among the selected local experts ----- + int32_t const localBegin = params.mLocalExpertsStartIdx; + int32_t const localEnd = params.mLocalExpertsStartIdx + params.mNumLocalExperts; + uint4 const lo = *reinterpret_cast(&smemSelected[0]); + uint4 const hi = *reinterpret_cast(&smemSelected[4]); + uint32_t selected[8] = {lo.x, lo.y, lo.z, lo.w, hi.x, hi.y, hi.z, hi.w}; + if (localBegin != 0 || localEnd < kNumExperts) + { +#pragma unroll + for (int i = 0; i < 8; ++i) + { + int32_t const b0 = max(localBegin - 32 * i, 0); + int32_t const b1 = min(localEnd - 32 * i, 32); + uint32_t const upTo = b1 >= 32 ? 0xFFFFFFFFu : ((1u << max(b1, 0)) - 1u); + uint32_t const below = b0 >= 32 ? 0xFFFFFFFFu : ((1u << b0) - 1u); + selected[i] &= b1 <= b0 ? 0u : (upTo & ~below); + } + } + int32_t numSelected = 0; +#pragma unroll + for (int i = 0; i < 8; ++i) + { + numSelected += __popc(selected[i]); + } + // (rank, is selected and local) of an expert, from masked popcounts of every word: indexing the words by the expert + // would place them in local memory. + auto rankOf = [&](int32_t e, bool& isSelected) + { + int32_t const word = e >> 5; + uint32_t const bit = 1u << (e & 31); + int32_t rank = 0; + int32_t hit = 0; +#pragma unroll + for (int i = 0; i < 8; ++i) + { + rank += __popc(selected[i] & (i < word ? 0xFFFFFFFFu : (i == word ? bit - 1u : 0u))); + hit += __popc(selected[i] & (i == word ? bit : 0u)); + } + isSelected = hit != 0; + return rank; + }; + auto tileRows = [&](int32_t n) + { return params.mIsPow2 ? mulLog2(n, params.mPaddingLog2) : mulTileN(n, params.mTileTokensDim); }; + + { + int32_t const expert = threadIdx.x; + bool isSelected; + int32_t const rank = rankOf(expert, isSelected); + if (isSelected) + { + uint2 const kRow = *reinterpret_cast(&smemKIdx[expert * kKIdxRowBytes]); + int32_t const count = __popc(__vsetne4(kRow.x, 0xFFFFFFFFu)) + __popc(__vsetne4(kRow.y, 0xFFFFFFFFu)); + params.mPtrCtaIdxXyToBatchIdx[rank] = expert - localBegin; + params.mPtrCtaIdxXyToMnLimit[rank] = min(tileRows(rank + 1), tileRows(rank) + count); + } + } + if (threadIdx.x == 0) + { + params.mPtrPermutedIdxSize[0] = tileRows(numSelected); + params.mPtrNumNonExitingCtas[0] = numSelected; + } + + // Permutation, one (token, k) pair per thread. offset-within-expert = number of earlier tokens that also selected + // the expert (its byte-map row below byte t), the small-token kernel's order. + if (static_cast(threadIdx.x) < kNumTokens * kTopK) + { + int32_t const expandedIdx = threadIdx.x; + int32_t const tokenIdx = expandedIdx / kTopK; + int32_t const pairExpert = smemTopKExpert[expandedIdx]; + bool pairIsLocal; + int32_t const rank = rankOf(pairExpert, pairIsLocal); + uint2 const row = *reinterpret_cast(&smemKIdx[pairExpert * kKIdxRowBytes]); + uint64_t const present = (static_cast(__vsetne4(row.y, 0xFFFFFFFFu)) << 32) + | static_cast(__vsetne4(row.x, 0xFFFFFFFFu)); + uint64_t const earlier = tokenIdx == 0 ? 0ull : (present & ((1ull << (8 * tokenIdx)) - 1)); + int32_t const permutedIdx = pairIsLocal ? tileRows(rank) + __popcll(earlier) : int32_t{-1}; + if (params.mPtrExpandedIdxToPermutedIdx != nullptr) + { + params.mPtrExpandedIdxToPermutedIdx[expandedIdx] = permutedIdx; + } + if (params.mPtrPermutedIdxToExpandedIdx != nullptr && pairIsLocal) + { + params.mPtrPermutedIdxToExpandedIdx[permutedIdx] = expandedIdx; + } + if (params.mPtrPermutedIdxToTokenIdx != nullptr && pairIsLocal) + { + params.mPtrPermutedIdxToTokenIdx[permutedIdx] = tokenIdx; + } + } + if (threadIdx.x == 0) + { + *epochPtr = epoch + 1u; + } + // After all routing outputs are written. + cudaTriggerProgrammaticLaunchCompletion(); +} +#else +__global__ void __launch_bounds__(kNumThreads, 1) + routingFusedRouterKernel(FusedRouterParams const __grid_constant__ /* p */) +{ + assert(false && "routingFusedRouterKernel is only supported on SM100+ architectures"); +} +#endif // if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) + +// --------------------------------------------------------------------------------------------------------------------- +// 9-24 tokens: a router GEMV front end followed in the same grid by a routing cluster. +// +// The front ends compute the logits of the production router kernel at that token count, bit for bit: +// - 9-16 tokens: a GEMV whose logits are router_gemm_kernel's (glmRouterGemm.cu): clusters +// of 4 CTAs of 256 threads, one per K slice of 8 experts; the final sum is added to 0.0f in slice order. +// - 17-24 tokens: router_gemm_mma_kernel's tensor-core GEMM (glmRouterGemm.cu): clusters of 8 +// CTAs of 128 threads, one per K slice of 768 for 16 experts; rank 0 adds the 8 slices in rank order. +// Both load their weights before the grid-dependency wait. Instead of a logits tensor, each cluster's rank 0 stores the +// logits into the exchange buffer [token][expert] of the workspace (by the parity of the epoch word), whose idle value +// is -0.0. A -0.0 logit is stored as +0.0, which gives the same score (sigmoid of either is 0.5); the sliced front end +// never produces one, since its final sum starts from 0.0f. +// +// Routing cluster (the grid's last cluster): warp w of rank r >= 1 routes token W (r - 1) + w (W warps per CTA). It +// polls the token's 256 logits until none is -0.0 and runs routingIndicesSmallTokenKernel's phase A (sigmoid + bias, +// reduceTopK, scaled sum normalization, with the bias staged in shared memory), which computes the values of the +// routing kernels' TopK. It writes the top-K weights and stores the top-K experts into rank 0's shared memory, whose +// slots rank 0 marked empty before the cluster's first barrier phase. Once every slot is filled, rank 0 writes +// routingIndicesDynBlockKernel's metadata: an expert with n local tokens owns ceil(n / tile) CTA tiles, the tiles +// follow expert order, and an expert's rows follow token order (the cluster routing kernel used at 17+ tokens writes +// the same tiles, with the rows of an expert in the order of its shared-memory atomics). Both the token warps and +// rank 0 first run their code once on dummy data in shared memory and registers, before the grid-dependency wait +// (while the primary grid runs): the live pass then executes from a warm instruction cache. The cluster also resets +// the previous launch's buffer and rank 0 advances the epoch, as the 1-8 token kernel does; both kernels keep the +// workspace in the same state. +namespace wide +{ +// Sliced front end, 9-16 tokens. +constexpr int kSliceMinTokens = 9; +constexpr int kSliceMaxTokens = 16; +constexpr int kSliceCluster = 4; // K slices of the 9-16-token GEMV +constexpr int kSliceThreads = 256; +constexpr int kSliceStageGroups = 3; // cp.async groups of 2 k-iterations +constexpr int kSliceKItersPerGroup = kKIters / kSliceStageGroups; +constexpr int kSliceRowBytes = 16; + +__host__ __device__ constexpr int slicePaddedTokens(int numTokens) +{ + return (numTokens + 3) / 4 * 4; +} + +// GEMV CTAs: activations [tokens][kKIters][32] 16-byte vectors, reused for the butterfly staging [32][8][paddedTokens] +// fp32, then the slice sums [kSliceCluster][8][tokens] fp32 received by rank 0 (router_gemm_sliced_kernel's layout). +__host__ __device__ constexpr int sliceWorkBytes(int numTokens) +{ + return numTokens * kKIters * WarpSize * 16 > WarpSize * kExpertsPerCta * slicePaddedTokens(numTokens) * 4 + ? numTokens * kKIters * WarpSize * 16 + : WarpSize * kExpertsPerCta * slicePaddedTokens(numTokens) * 4; +} + +__host__ __device__ constexpr int sliceGemvSmemBytes(int numTokens) +{ + return sliceWorkBytes(numTokens) + kSliceCluster * kExpertsPerCta * numTokens * 4; +} + +// Tensor-core front end, 17-24 tokens (router_gemm_mma_kernel's routerMma constants). +constexpr int kMmaMinTokens = 17; +constexpr int kMmaMaxTokens = 24; +constexpr int kMmaExpertsPerCta = 16; +constexpr int kMmaCluster = 8; // K slices of 768 +constexpr int kMmaSliceK = kHiddenDim / kMmaCluster; +constexpr int kMmaRowElems = kMmaSliceK + 8; // the 8 rows of an ldmatrix start in distinct banks +constexpr int kMmaWarps = 4; +constexpr int kMmaThreads = kMmaWarps * WarpSize; +constexpr int kMmaKStepsPerWarp = kMmaSliceK / 16 / kMmaWarps; +constexpr int kMmaChunks = kMmaSliceK / 8; // 16-byte chunks per row of a slice +constexpr int kMmaGemvCtas = kNumExperts / kMmaExpertsPerCta * kMmaCluster; +constexpr int kMmaRowBytes = 32; + +__host__ __device__ constexpr int mmaPaddedTokens(int numTokens) +{ + return (numTokens + 7) / 8 * 8; +} + +// Weight tile and activation rows (bf16), the 4 warp tiles and the 8 slice tiles received by rank 0 (fp32). +__host__ __device__ constexpr int mmaGemvSmemBytes(int numTokens) +{ + return (kMmaExpertsPerCta + mmaPaddedTokens(numTokens)) * kMmaRowElems * 2 + + (kMmaWarps + kMmaCluster) * kMmaExpertsPerCta * mmaPaddedTokens(numTokens) * 4; +} + +// Routing CTAs: bias [256] fp32, top-K experts [kTokenCapacity][8] int32, byte map [256][kRowBytes] (one byte per +// token, 0xFF = not selected), CTA offsets [256] int32, warp totals [8] int32. +template +struct RoutingSmem +{ + static constexpr int kBias = 0; + static constexpr int kTopKExperts = kBias + kNumExperts * 4; + static constexpr int kKIdx = kTopKExperts + kTokenCapacity * kTopK * 4; + static constexpr int kCtaOffset = kKIdx + kNumExperts * kRowBytes; + static constexpr int kWarpTotals = kCtaOffset + kNumExperts * 4; + static constexpr int kBytes = kWarpTotals + 8 * 4; + static_assert(kKIdx % 16 == 0, "byte map rows are read as 16-byte vectors"); +}; + +using SliceRoutingSmem = RoutingSmem; +using MmaRoutingSmem = RoutingSmem; + +__host__ __device__ constexpr int sliceSmemBytes(int numTokens) +{ + return sliceGemvSmemBytes(numTokens) > SliceRoutingSmem::kBytes ? sliceGemvSmemBytes(numTokens) + : SliceRoutingSmem::kBytes; +} + +__host__ __device__ constexpr int mmaSmemBytes(int numTokens) +{ + return mmaGemvSmemBytes(numTokens) > MmaRoutingSmem::kBytes ? mmaGemvSmemBytes(numTokens) : MmaRoutingSmem::kBytes; +} + +static_assert(kMmaMaxTokens * kNumExperts <= kExchangeFloats, "a buffer holds [kMmaMaxTokens][256] logits"); +static_assert(kSliceCluster == kSlices && kNumGemvCtas == kNumExperts / kExpertsPerCta * kSliceCluster, "GEMV grid"); +} // namespace wide + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +__device__ __forceinline__ void clusterArriveRelaxed() +{ + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); +} + +__device__ __forceinline__ void clusterWait() +{ + asm volatile("barrier.cluster.wait.aligned;" ::: "memory"); +} + +__device__ __forceinline__ void clusterArriveRelease() +{ + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); +} + +__device__ __forceinline__ void clusterWaitAcquire() +{ + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); +} + +__device__ __forceinline__ uint32_t mapToClusterRank(void const* ptr, uint32_t rank) +{ + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(smemAddr(ptr)), "r"(rank)); + return remote; +} + +// Stores v at the same shared-memory offset as ptr in the CTA of the given rank of this cluster. +__device__ __forceinline__ void storeToClusterRank(float* ptr, uint32_t rank, float v) +{ + asm volatile("st.shared::cluster.f32 [%0], %1;" ::"r"(mapToClusterRank(ptr, rank)), "f"(v) : "memory"); +} + +__device__ __forceinline__ void storeRelaxedToClusterRank(int32_t* ptr, uint32_t rank, int32_t v) +{ + asm volatile("st.relaxed.cluster.shared::cluster.s32 [%0], %1;" ::"r"(mapToClusterRank(ptr, rank)), "r"(v) + : "memory"); +} + +// A slot of this CTA's shared memory that other CTAs of the cluster store to. +__device__ __forceinline__ int32_t loadRelaxedCluster(int32_t const* ptr) +{ + int32_t v; + asm volatile("ld.relaxed.cluster.shared::cta.s32 %0, [%1];" : "=r"(v) : "r"(smemAddr(ptr)) : "memory"); + return v; +} + +__device__ __forceinline__ void copyAsync16(void* dst, void const* src) +{ + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ::"r"(smemAddr(dst)), "l"(src) : "memory"); +} + +template +__device__ __forceinline__ void copyAsyncWait() +{ + asm volatile("cp.async.wait_group %0;" ::"n"(kPending) : "memory"); +} + +__device__ __forceinline__ void ldmatrixX4(uint32_t (&r)[4], void const* ptr) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(smemAddr(ptr))); +} + +__device__ __forceinline__ void ldmatrixX2(uint32_t (&r)[2], void const* ptr) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];" + : "=r"(r[0]), "=r"(r[1]) + : "r"(smemAddr(ptr))); +} + +__device__ __forceinline__ void mmaBf16(float (&c)[4], uint32_t const (&a)[4], uint32_t const (&b)[2]) +{ + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); +} + +// The number of tokens before token beforeToken (all tokens if beforeToken >= kRowBytes) in an expert's byte map row. +template +__device__ __forceinline__ int32_t countTokensBefore(uint8_t const* smemKIdx, int32_t expert, int32_t beforeToken) +{ + uint4 const* row = reinterpret_cast(smemKIdx + expert * kRowBytes); + int32_t n = 0; +#pragma unroll + for (int v = 0; v < kRowBytes / 16; ++v) + { + uint4 const w4 = row[v]; + uint32_t const words[4] = {w4.x, w4.y, w4.z, w4.w}; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + int32_t const below = beforeToken - (16 * v + 4 * q); // tokens of this word before beforeToken + uint32_t const mask = below >= 4 ? 0xFFFFFFFFu : (below <= 0 ? 0u : (1u << (8 * below)) - 1u); + n += __popc(__vsetne4(words[q], 0xFFFFFFFFu) & mask); + } + } + return n; +} + +#endif + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) +// The routing cluster of the 9-24 token kernels (see above): kThreads per CTA, kClusterSize CTAs. +template +__device__ __forceinline__ void routingClusterTail(FusedRouterParams const& p, uint8_t* smem, uint32_t rank, + int32_t warpIdx, int32_t laneIdx, float* buffers, uint32_t* epochPtr) +{ + using Layout = wide::RoutingSmem; + constexpr int kWarpsPerCta = kThreads / WarpSize; + constexpr int kExpertsPerThread = kNumExperts / kThreads; + static_assert(kNumExperts % kThreads == 0, "whole experts per thread in rank 0"); + static_assert(kNumTokens <= kTokenCapacity && kTokenCapacity <= kRowBytes && kRowBytes % 16 == 0, "byte map rows"); + static_assert((kClusterSize - 1) * kWarpsPerCta >= kNumTokens, "one warp per token in ranks 1 and up"); + static_assert(kWarpsPerCta <= 8, "warp totals"); + auto const& params = p.routing; + float* const smemBias = reinterpret_cast(smem + Layout::kBias); + int32_t* const smemTopKExpert = reinterpret_cast(smem + Layout::kTopKExperts); + uint8_t* const smemKIdx = smem + Layout::kKIdx; + int32_t* const smemCtaOffset = reinterpret_cast(smem + Layout::kCtaOffset); + int32_t* const smemWarpTotals = reinterpret_cast(smem + Layout::kWarpTotals); + // Independent of the primary grid: stage the routing bias (the value the preprocess computes per element) and, in + // rank 0, clear the byte map and mark the top-K slots empty (-1). + auto const& pre = params.mExpertSelectParams.mPreprocessParams; + for (int e = threadIdx.x; e < kNumExperts; e += kThreads) + { + smemBias[e] = loadScalar(pre.ptrRoutingBias, e, pre.dtypeBias); + } + if (rank == 0) + { + for (int i = threadIdx.x; i < kNumExperts * kRowBytes / 16; i += kThreads) + { + reinterpret_cast(smemKIdx)[i] = make_uint4(0xFFFFFFFFu, 0xFFFFFFFFu, 0xFFFFFFFFu, 0xFFFFFFFFu); + } + for (int i = threadIdx.x; i < kNumTokens * kTopK; i += kThreads) + { + smemTopKExpert[i] = -1; + } + } + __syncthreads(); + // First phase of the cluster barrier, after rank 0 has marked the top-K slots empty: the other ranks store to them + // only after the wait. + clusterArriveRelaxed(); + // Each loop below runs a pass 0 on dummy data in shared memory and registers only, before the grid-dependency + // wait (it overlaps the primary grid), so that pass 1 runs from a warm instruction cache. Pass 1 starts with the + // wait, the epoch, the reset of the previous launch's buffer and the wait of the cluster barrier's first phase. + uint32_t epoch = 0; + float const* exchange = nullptr; + auto const enterLivePass = [&]() + { + gridDependencyWait(); + epoch = loadVolatile(epochPtr); + exchange = buffers + (epoch & 1u) * kExchangeFloats; + // The buffer of the previous launch, which its routing cluster has consumed: idle again for the next launch. + float4* other = reinterpret_cast(buffers + ((epoch + 1u) & 1u) * kExchangeFloats); + for (int i = static_cast(rank) * kThreads + threadIdx.x; i < kExchangeFloats / 4; + i += kClusterSize * kThreads) + { + other[i] = make_float4(-0.f, -0.f, -0.f, -0.f); + } + clusterWait(); + }; + + if (rank != 0) + { + // ----- Phase A: one warp per token, as in routingIndicesSmallTokenKernel ----- + // Lane l holds experts 8 l .. 8 l + 7 (2 16-byte loads); the top-K is independent of how the experts are spread + // over the lanes. + int32_t const tokenIdx = (static_cast(rank) - 1) * kWarpsPerCta + warpIdx; + bool const hasToken = tokenIdx < kNumTokens; + auto block = cg::this_thread_block(); + auto warp = cg::tiled_partition(block); + constexpr int VecSize = kNumExperts / WarpSize; + static_assert(VecSize == 8, "two 16-byte loads per lane"); +#pragma unroll 1 + for (int pass = 0; pass < 2; ++pass) + { + bool const live = pass == 1; + if (live) + { + enterLivePass(); + } + if (hasToken) + { + float score[VecSize]; + int32_t idx[VecSize]; + if (live) + { + float const* const row = exchange + tokenIdx * kNumExperts + laneIdx * VecSize; + bool ready; + do + { + float4 const lo = loadVolatile4(row); + float4 const hi = loadVolatile4(row + 4); + score[0] = lo.x; + score[1] = lo.y; + score[2] = lo.z; + score[3] = lo.w; + score[4] = hi.x; + score[5] = hi.y; + score[6] = hi.z; + score[7] = hi.w; + ready = true; +#pragma unroll + for (int i = 0; i < VecSize; ++i) + { + ready = ready && !isNegZero(score[i]); + } + } while (!ready); + } + else + { +#pragma unroll + for (int i = 0; i < VecSize; ++i) + { + score[i] = 0.f; + } + } + // SigmoidBiasPreprocess::apply + reduceTopK + ScaledSumNormalizePostprocess::apply, with the bias read + // from shared memory: the same operations on the same values in the same order. +#pragma unroll + for (int i = 0; i < VecSize; ++i) + { + idx[i] = laneIdx * VecSize + i; + float const s = sigmoid_accurate(score[i]); + score[i] = s + smemBias[idx[i]]; + } + float topScore[kTopK]; + int32_t topExpert[kTopK]; + topk::reduceTopK(warp, topScore, topExpert, score, idx, float{-INFINITY}, kTopK); + // Lane k's top-K entry. + float myScore = 0.f; + int32_t myExpert = 0; +#pragma unroll + for (int k = 0; k < kTopK; ++k) + { + myScore = laneIdx == k ? topScore[k] : myScore; + myExpert = laneIdx == k ? topExpert[k] : myExpert; + } + bool const isTop = laneIdx < kTopK; + auto const& post = params.mExpertSelectParams.mPostprocessParams; + float const biasVal = isTop ? smemBias[myExpert] : 0.f; + float const sigmoidScore = isTop ? (myScore - biasVal) : 0.f; + float const sum = cg::reduce(warp, sigmoidScore, cg::plus()); + if (live && isTop) + { + float const weight = sigmoidScore * post.routeScale / (sum + post.sumEpsilon); + params.mPtrTopKWeights[tokenIdx * kTopK + laneIdx] = __nv_bfloat16{weight}; + storeRelaxedToClusterRank(&smemTopKExpert[tokenIdx * kTopK + laneIdx], 0, myExpert); + } + } + } + // Rank 0 polls its top-K slots; this CTA is done. + return; + } + + // ----- Metadata (rank 0), as routingIndicesDynBlockKernel ----- + // Pass 0 runs on the cleared byte map and expert 0, without global stores. + int32_t const localBegin = params.mLocalExpertsStartIdx; + int32_t const localEnd = params.mLocalExpertsStartIdx + params.mNumLocalExperts; + auto const tileRows = [&](int32_t n) + { return params.mIsPow2 ? mulLog2(n, params.mPaddingLog2) : mulTileN(n, params.mTileTokensDim); }; +#pragma unroll 1 + for (int pass = 0; pass < 2; ++pass) + { + bool const live = pass == 1; + if (live) + { + enterLivePass(); + // Every (token, k) slot is stored once by the token's warp; wait for it. + for (int pr = threadIdx.x; pr < kNumTokens * kTopK; pr += kThreads) + { + int32_t expert; + do + { + expert = loadRelaxedCluster(&smemTopKExpert[pr]); + } while (expert < 0); + smemKIdx[expert * kRowBytes + pr / kTopK] = static_cast(pr % kTopK); + } + } + __syncthreads(); + // Thread t owns experts kExpertsPerThread t + j; an expert's CTA tiles follow those of the experts before it. + int32_t count[kExpertsPerThread]; + int32_t numCta[kExpertsPerThread]; + int32_t threadCtas = 0; +#pragma unroll + for (int j = 0; j < kExpertsPerThread; ++j) + { + int32_t const expert = threadIdx.x * kExpertsPerThread + j; + count[j] = countTokensBefore(smemKIdx, expert, kRowBytes); + bool const isLocal = expert >= localBegin && expert < localEnd; + numCta[j] = !isLocal ? 0 + : (params.mIsPow2 ? divUpLog2(count[j], params.mPaddingLog2) + : divUpTileN(count[j], params.mTileTokensDim)); + threadCtas += numCta[j]; + } + int32_t inclusive = threadCtas; +#pragma unroll + for (int offset = 1; offset < WarpSize; offset *= 2) + { + int32_t const v = __shfl_up_sync(0xffffffff, inclusive, offset); + inclusive += laneIdx >= offset ? v : 0; + } + if (laneIdx == WarpSize - 1) + { + smemWarpTotals[warpIdx] = inclusive; + } + __syncthreads(); + int32_t warpPrefix = 0; + int32_t numNonExitingCtas = 0; +#pragma unroll + for (int w = 0; w < kWarpsPerCta; ++w) + { + int32_t const total = smemWarpTotals[w]; + warpPrefix += w < warpIdx ? total : 0; + numNonExitingCtas += total; + } + int32_t ctaOffset = warpPrefix + inclusive - threadCtas; +#pragma unroll + for (int j = 0; j < kExpertsPerThread; ++j) + { + int32_t const expert = threadIdx.x * kExpertsPerThread + j; + smemCtaOffset[expert] = ctaOffset; + for (int32_t cta = 0; cta < numCta[j]; ++cta) + { + params.mPtrCtaIdxXyToBatchIdx[ctaOffset + cta] = expert - localBegin; + params.mPtrCtaIdxXyToMnLimit[ctaOffset + cta] + = min(tileRows(ctaOffset + cta + 1), tileRows(ctaOffset) + count[j]); + } + ctaOffset += numCta[j]; + } + if (live && threadIdx.x == 0) + { + params.mPtrPermutedIdxSize[0] = tileRows(numNonExitingCtas); + params.mPtrNumNonExitingCtas[0] = numNonExitingCtas; + } + __syncthreads(); + + // Permutation, one (token, k) pair per thread and pass of the loop: rows of an expert in token order. + for (int pr = threadIdx.x; pr < kNumTokens * kTopK; pr += kThreads) + { + int32_t const expandedIdx = pr; + int32_t const pairToken = expandedIdx / kTopK; + // The warm-up pass reads expert 0's (cleared) row: smemTopKExpert is written by the other ranks. + int32_t const pairExpert = live ? smemTopKExpert[expandedIdx] : 0; + bool const pairIsLocal = pairExpert >= localBegin && pairExpert < localEnd; + int32_t const earlier = countTokensBefore(smemKIdx, pairExpert, pairToken); + int32_t const permutedIdx = pairIsLocal ? tileRows(smemCtaOffset[pairExpert]) + earlier : int32_t{-1}; + if (live && params.mPtrExpandedIdxToPermutedIdx != nullptr) + { + params.mPtrExpandedIdxToPermutedIdx[expandedIdx] = permutedIdx; + } + if (live && params.mPtrPermutedIdxToExpandedIdx != nullptr && pairIsLocal) + { + params.mPtrPermutedIdxToExpandedIdx[permutedIdx] = expandedIdx; + } + if (live && params.mPtrPermutedIdxToTokenIdx != nullptr && pairIsLocal) + { + params.mPtrPermutedIdxToTokenIdx[permutedIdx] = pairToken; + } + } + } + if (threadIdx.x == 0) + { + *epochPtr = epoch + 1u; + } + // After all routing outputs are written. + cudaTriggerProgrammaticLaunchCompletion(); +} +#endif // if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) + +template +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) +__global__ void __launch_bounds__(wide::kSliceThreads, 1) + routingFusedRouterWideKernel(FusedRouterParams const __grid_constant__ p) +{ + using namespace wide; + extern __shared__ __align__(128) uint8_t smemWide[]; + uint8_t* const smem = smemWide; + float* const buffers = reinterpret_cast(p.workspace); + uint32_t* const epochPtr = reinterpret_cast(p.workspace + 2 * kExchangeFloats); + uint32_t const rank = blockIdx.x % kSliceCluster; + int32_t const warpIdx = __shfl_sync(0xffffffff, threadIdx.x / WarpSize, 0); + int32_t const laneIdx = cutlass::arch::LaneId(); + + if (blockIdx.x < kNumGemvCtas) + { + // First phase of the cluster barrier: arrive now, wait before the first remote store, so that every CTA of the + // cluster has started by then. + clusterArriveRelaxed(); + constexpr int kPadded = slicePaddedTokens(kNumTokens); + constexpr int kPairs = kExpertsPerCta * kNumTokens; + uint4* act = reinterpret_cast(smem); + float* fold = reinterpret_cast(smem); + float* sliceSums = reinterpret_cast(smem + sliceWorkBytes(kNumTokens)); + int const slice = static_cast(rank); + int const expertBase = blockIdx.x / kSliceCluster * kExpertsPerCta; + // Position t of the slice (router_gemm_kernel's lane) and expert of this thread; the 4 experts of a warp read + // the same 128 bytes of activations. + int const t = 8 * (warpIdx % 4) + laneIdx % 8; + int const eLocal = 4 * (warpIdx / 4) + laneIdx / 8; + __nv_bfloat16 const* wRow = p.routerWeight + (expertBase + eLocal) * kHiddenDim + (WarpSize * slice + t) * 8; + // The weights are not written by the primary grid: load them before the wait. A shared-memory store of their + // XOR ahead of a barrier keeps the loads before the wait. + uint4 w[kKIters]; + uint32_t loaded = 0; +#pragma unroll + for (int ki = 0; ki < kKIters; ki++) + { + w[ki] = loadWeight16(wRow + ki * kKIterElts); + loaded ^= w[ki].x ^ w[ki].y ^ w[ki].z ^ w[ki].w; + } + reinterpret_cast(fold)[threadIdx.x] = loaded; + __syncthreads(); + gridDependencyWait(); + // The previous launch on this workspace precedes the primary grid in the stream, so it has completed. + uint32_t const epoch = slice == 0 ? loadVolatile(epochPtr) : 0u; + // Stage the slice: vector (m, ki, pos) holds token m's elements 1024 ki + 8 (32 slice + pos) + 0..7. +#pragma unroll + for (int group = 0; group < kSliceStageGroups; group++) + { + constexpr int kVecs = kNumTokens * kSliceKItersPerGroup * WarpSize; + for (int c = threadIdx.x; c < kVecs; c += kSliceThreads) + { + int const m = c / (kSliceKItersPerGroup * WarpSize); + int const rem = c - m * (kSliceKItersPerGroup * WarpSize); + int const ki = group * kSliceKItersPerGroup + rem / WarpSize; + int const pos = rem % WarpSize; + copyAsync16(&act[(m * kKIters + ki) * WarpSize + pos], + p.hiddenStates + static_cast(m) * kHiddenDim + ki * kKIterElts + + (WarpSize * slice + pos) * 8); + } + asm volatile("cp.async.commit_group;" ::: "memory"); + } + // The per-thread chains of router_gemm_kernel's thread 32 slice + t, for one expert and every token. + float acc[kPadded]; +#pragma unroll + for (int m = 0; m < kPadded; m++) + { + acc[m] = 0.f; + } +#pragma unroll + for (int group = 0; group < kSliceStageGroups; group++) + { + if (group == 0) + { + copyAsyncWait(); + } + else if (group == 1) + { + copyAsyncWait(); + } + else + { + copyAsyncWait<0>(); + } + __syncthreads(); +#pragma unroll + for (int kk = 0; kk < kSliceKItersPerGroup; kk++) + { + int const ki = group * kSliceKItersPerGroup + kk; +#pragma unroll + for (int m = 0; m < kNumTokens; m++) + { + fmaBf16x8(acc[m], act[(m * kKIters + ki) * WarpSize + t], w[ki]); + } + } + } + // Every thread is done with the activations; the butterfly staging reuses their space. + __syncthreads(); + { + float4* row = reinterpret_cast(fold + (t * kExpertsPerCta + eLocal) * kPadded); +#pragma unroll + for (int q = 0; q < kPadded / 4; q++) + { + row[q] = make_float4(acc[4 * q], acc[4 * q + 1], acc[4 * q + 2], acc[4 * q + 3]); + } + } + __syncthreads(); + clusterWait(); + // The butterfly of router_gemm_kernel's warp: 4 threads per (expert, token) fold the positions + // {sub, sub + 4, ..., sub + 28} through levels 16, 8 and 4 from shared memory, two shuffles add levels 2, 1. +#pragma unroll + for (int base = 0; base < kPairs; base += kSliceThreads / 4) + { + int const pairRaw = base + threadIdx.x / 4; + bool const valid = pairRaw < kPairs; + int const pair = valid ? pairRaw : 0; + int const sub = threadIdx.x % 4; + int const eL = pair / kNumTokens; + int const m = pair % kNumTokens; + auto const v = [&](int position) { return fold[(position * kExpertsPerCta + eL) * kPadded + m]; }; + float const s1a = v(sub) + v(sub + 16); + float const s1b = v(sub + 4) + v(sub + 20); + float const s1c = v(sub + 8) + v(sub + 24); + float const s1d = v(sub + 12) + v(sub + 28); + float sum = (s1a + s1c) + (s1b + s1d); + sum += __shfl_xor_sync(0xffffffff, sum, 2); + sum += __shfl_xor_sync(0xffffffff, sum, 1); + if (valid && sub == 0) + { + storeToClusterRank(&sliceSums[slice * kPairs + pair], 0, sum); + } + } + // Second phase: rank 0 waits for the slice sums; the other CTAs are done (nothing reads their shared memory). + clusterArriveRelease(); + if (slice != 0) + { + return; + } + clusterWaitAcquire(); + float* const exchange = buffers + (epoch & 1u) * kExchangeFloats; + for (int pr = threadIdx.x; pr < kPairs; pr += kSliceThreads) + { + float finalSum = 0.0f; +#pragma unroll + for (int s = 0; s < kSliceCluster; s++) + { + finalSum += sliceSums[s * kPairs + pr]; + } + storeRelaxed(&exchange[(pr % kNumTokens) * kNumExperts + expertBase + pr / kNumTokens], finalSum); + } + // Exit without a trigger: the dependents launch once the routing cluster's rank 0 triggers. + return; + } + routingClusterTail( + p, smem, rank, warpIdx, laneIdx, buffers, epochPtr); +} +#else +__global__ void __launch_bounds__(wide::kSliceThreads, 1) + routingFusedRouterWideKernel(FusedRouterParams const __grid_constant__ /* p */) +{ + assert(false && "routingFusedRouterWideKernel is only supported on SM100+ architectures"); +} +#endif // if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) + +template +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) +__global__ void __launch_bounds__(wide::kMmaThreads, 1) + routingFusedRouterMmaKernel(FusedRouterParams const __grid_constant__ p) +{ + using namespace wide; + extern __shared__ __align__(128) uint8_t smemMmaFused[]; + uint8_t* const smem = smemMmaFused; + float* const buffers = reinterpret_cast(p.workspace); + uint32_t* const epochPtr = reinterpret_cast(p.workspace + 2 * kExchangeFloats); + uint32_t const rank = blockIdx.x % kMmaCluster; + int32_t const warpIdx = __shfl_sync(0xffffffff, threadIdx.x / WarpSize, 0); + int32_t const laneIdx = cutlass::arch::LaneId(); + + if (blockIdx.x < kMmaGemvCtas) + { + // First phase of the cluster barrier: arrive now, wait before the first remote store, so that every CTA of the + // cluster has started by then. + clusterArriveRelaxed(); + // router_gemm_mma_kernel, with the weight tile copied before the wait. + constexpr int kPadded = mmaPaddedTokens(kNumTokens); + constexpr int kTokenTiles = kPadded / 8; + constexpr int kTile = kMmaExpertsPerCta * kPadded; + // The weight tile [16][kMmaRowElems] and the activation rows [kPadded][kMmaRowElems] (bf16), the warp tiles + // [kMmaWarps][16][kPadded] and, used in rank 0, the slice tiles [kMmaCluster][16][kPadded] (fp32). + __nv_bfloat16* weightTile = reinterpret_cast<__nv_bfloat16*>(smem); + __nv_bfloat16* actTile = weightTile + kMmaExpertsPerCta * kMmaRowElems; + float* warpSums = reinterpret_cast(actTile + kPadded * kMmaRowElems); + float* sliceSums = warpSums + kMmaWarps * kTile; + int const slice = static_cast(rank); + int const expertBase = blockIdx.x / kMmaCluster * kMmaExpertsPerCta; + // The primary grid does not write the weights. + for (int c = threadIdx.x; c < kMmaExpertsPerCta * kMmaChunks; c += kMmaThreads) + { + int const row = c / kMmaChunks; + int const chunk = c % kMmaChunks; + copyAsync16(weightTile + row * kMmaRowElems + chunk * 8, + p.routerWeight + static_cast(expertBase + row) * kHiddenDim + slice * kMmaSliceK + chunk * 8); + } + asm volatile("cp.async.commit_group;" ::: "memory"); + // Rows past the last token are zero; they only feed outputs that are not stored. + for (int c = threadIdx.x; c < (kPadded - kNumTokens) * kMmaChunks; c += kMmaThreads) + { + *reinterpret_cast(actTile + (kNumTokens + c / kMmaChunks) * kMmaRowElems + (c % kMmaChunks) * 8) + = make_uint4(0, 0, 0, 0); + } + gridDependencyWait(); + // The previous launch on this workspace precedes the primary grid in the stream, so it has completed. + uint32_t const epoch = slice == 0 ? loadVolatile(epochPtr) : 0u; + for (int c = threadIdx.x; c < kNumTokens * kMmaChunks; c += kMmaThreads) + { + int const row = c / kMmaChunks; + int const chunk = c % kMmaChunks; + copyAsync16(actTile + row * kMmaRowElems + chunk * 8, + p.hiddenStates + static_cast(row) * kHiddenDim + slice * kMmaSliceK + chunk * 8); + } + asm volatile("cp.async.commit_group;" ::: "memory"); + copyAsyncWait<0>(); + __syncthreads(); + + float acc[kTokenTiles][4] = {}; +#pragma unroll + for (int s = 0; s < kMmaKStepsPerWarp; s++) + { + int const k0 = (warpIdx * kMmaKStepsPerWarp + s) * 16; + uint32_t a[4]; + ldmatrixX4(a, weightTile + (laneIdx % 16) * kMmaRowElems + k0 + (laneIdx / 16) * 8); +#pragma unroll + for (int j = 0; j < kTokenTiles; j++) + { + uint32_t b[2]; + ldmatrixX2(b, actTile + (j * 8 + laneIdx % 8) * kMmaRowElems + k0 + ((laneIdx / 8) % 2) * 8); + mmaBf16(acc[j], a, b); + } + } + // m16n8 accumulators: acc[j][0..1] at expert lane / 4, tokens 8 j + 2 (lane % 4) + 0..1; acc[j][2..3] 8 experts + // on. +#pragma unroll + for (int j = 0; j < kTokenTiles; j++) + { + float* row = warpSums + (warpIdx * kMmaExpertsPerCta + laneIdx / 4) * kPadded + j * 8 + 2 * (laneIdx % 4); + *reinterpret_cast(row) = make_float2(acc[j][0], acc[j][1]); + *reinterpret_cast(row + 8 * kPadded) = make_float2(acc[j][2], acc[j][3]); + } + __syncthreads(); + clusterWait(); + for (int i = threadIdx.x; i < kTile; i += kMmaThreads) + { + if (i % kPadded < kNumTokens) + { + float const sum + = ((warpSums[i] + warpSums[kTile + i]) + warpSums[2 * kTile + i]) + warpSums[3 * kTile + i]; + storeToClusterRank(&sliceSums[slice * kTile + i], 0, sum); + } + } + // Second phase: rank 0 waits for the slices; the other CTAs are done (nothing reads their shared memory). + clusterArriveRelease(); + if (slice != 0) + { + return; + } + clusterWaitAcquire(); + float* const exchange = buffers + (epoch & 1u) * kExchangeFloats; + for (int pr = threadIdx.x; pr < kMmaExpertsPerCta * kNumTokens; pr += kMmaThreads) + { + int const e = pr / kNumTokens; + int const m = pr % kNumTokens; + float sum = sliceSums[e * kPadded + m]; +#pragma unroll + for (int r = 1; r < kMmaCluster; r++) + { + sum += sliceSums[r * kTile + e * kPadded + m]; + } + storeRelaxed(&exchange[m * kNumExperts + expertBase + e], isNegZero(sum) ? 0.f : sum); + } + // Exit without a trigger: the dependents launch once the routing cluster's rank 0 triggers. + return; + } + routingClusterTail( + p, smem, rank, warpIdx, laneIdx, buffers, epochPtr); +} +#else +__global__ void __launch_bounds__(wide::kMmaThreads, 1) + routingFusedRouterMmaKernel(FusedRouterParams const __grid_constant__ /* p */) +{ + assert(false && "routingFusedRouterMmaKernel is only supported on SM100+ architectures"); +} +#endif // if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) + +template +void launchClusterGrid(Kernel kernel, FusedRouterParams const& params, int gridDim, int blockDim, int clusterDim, + int smemBytes, bool usePdl, cudaStream_t stream) +{ + TLLM_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smemBytes)); + cudaLaunchConfig_t config{}; + config.gridDim = gridDim; + config.blockDim = blockDim; + config.dynamicSmemBytes = smemBytes; + config.stream = stream; + cudaLaunchAttribute attributes[2] = {}; + attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attributes[0].val.programmaticStreamSerializationAllowed = int(usePdl); + attributes[1].id = cudaLaunchAttributeClusterDimension; + attributes[1].val.clusterDim.x = clusterDim; + attributes[1].val.clusterDim.y = 1; + attributes[1].val.clusterDim.z = 1; + config.attrs = attributes; + config.numAttrs = 2; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, params)); +} + +template +void launchFusedRouterWide(FusedRouterParams const& params, bool usePdl, cudaStream_t stream) +{ + // 128 GEMV CTAs and one routing cluster. + launchClusterGrid(routingFusedRouterWideKernel, params, kNumGemvCtas + wide::kSliceCluster, + wide::kSliceThreads, wide::kSliceCluster, wide::sliceSmemBytes(kNumTokens), usePdl, stream); +} + +template +void launchFusedRouterMma(FusedRouterParams const& params, bool usePdl, cudaStream_t stream) +{ + // 128 GEMV CTAs and one routing cluster. + launchClusterGrid(routingFusedRouterMmaKernel, params, wide::kMmaGemvCtas + wide::kMmaCluster, + wide::kMmaThreads, wide::kMmaCluster, wide::mmaSmemBytes(kNumTokens), usePdl, stream); +} + +template +void launchFusedRouter(FusedRouterParams const& params, bool usePdl, cudaStream_t stream) +{ + cudaLaunchConfig_t config{}; + config.gridDim = kNumGemvCtas + 1; + config.blockDim = kNumThreads; + config.dynamicSmemBytes = kSmemBytes; + config.stream = stream; + cudaLaunchAttribute attributes[1] = {}; + attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attributes[0].val.programmaticStreamSerializationAllowed = int(usePdl); + config.attrs = attributes; + config.numAttrs = 1; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, routingFusedRouterKernel, params)); +} + +bool supportsFusedRouter(Data const& data, int32_t hiddenDim) +{ + // 1-8 tokens: at most one CTA tile per expert. 9-24 tokens: any tile size. + bool const tokensOk + = (data.mNumTokens >= 1 && data.mNumTokens <= kMaxTokens && data.mNumTokens <= data.mTileTokensDim) + || (data.mNumTokens >= wide::kSliceMinTokens && data.mNumTokens <= wide::kSliceMaxTokens) + || (data.mNumTokens >= wide::kMmaMinTokens && data.mNumTokens <= wide::kMmaMaxTokens); + return hiddenDim == kHiddenDim && data.mNumExperts == kNumExperts && data.mTopK == kTopK && tokensOk + && data.mLocalExpertsStrideLog2 == 0 && data.mPreprocessType == RoutingPreprocessType::SigmoidBias + && data.mPostprocessType == RoutingPostprocessType::ScaledSumNormalize + && data.mDtypeOutput == tg::Dtype::Bfloat16 && data.mPtrTopKWeights != nullptr + && data.mPtrRoutingBias != nullptr && data.mPtrTopKIds == nullptr; +} + +void runFusedRouter(Data const& data, void const* hiddenStates, void const* routerWeight, int32_t hiddenDim, + int32_t* workspace, void* stream) +{ + TLLM_CHECK_WITH_INFO(supportsFusedRouter(data, hiddenDim), + "The fused router + routing kernel supports 1-8 tokens (at most the tile size) or 9-24 tokens, 256 experts, " + "top-8, hidden size %d and SigmoidBias + ScaledSumNormalize routing; got %d tokens, tile %d, %d experts, " + "top-%d, hidden %d", + kHiddenDim, data.mNumTokens, data.mTileTokensDim, data.mNumExperts, data.mTopK, hiddenDim); + FusedRouterParams params; + params.routing = FusedKernelParams::setKernelParams(data); + params.hiddenStates = static_cast<__nv_bfloat16 const*>(hiddenStates); + params.routerWeight = static_cast<__nv_bfloat16 const*>(routerWeight); + params.workspace = workspace; + auto const s = static_cast(stream); + switch (data.mNumTokens) + { + case 1: launchFusedRouter<1>(params, data.mUsePdl, s); break; + case 2: launchFusedRouter<2>(params, data.mUsePdl, s); break; + case 3: launchFusedRouter<3>(params, data.mUsePdl, s); break; + case 4: launchFusedRouter<4>(params, data.mUsePdl, s); break; + case 5: launchFusedRouter<5>(params, data.mUsePdl, s); break; + case 6: launchFusedRouter<6>(params, data.mUsePdl, s); break; + case 7: launchFusedRouter<7>(params, data.mUsePdl, s); break; + case 8: launchFusedRouter<8>(params, data.mUsePdl, s); break; + case 9: launchFusedRouterWide<9>(params, data.mUsePdl, s); break; + case 10: launchFusedRouterWide<10>(params, data.mUsePdl, s); break; + case 11: launchFusedRouterWide<11>(params, data.mUsePdl, s); break; + case 12: launchFusedRouterWide<12>(params, data.mUsePdl, s); break; + case 13: launchFusedRouterWide<13>(params, data.mUsePdl, s); break; + case 14: launchFusedRouterWide<14>(params, data.mUsePdl, s); break; + case 15: launchFusedRouterWide<15>(params, data.mUsePdl, s); break; + case 16: launchFusedRouterWide<16>(params, data.mUsePdl, s); break; + case 17: launchFusedRouterMma<17>(params, data.mUsePdl, s); break; + case 18: launchFusedRouterMma<18>(params, data.mUsePdl, s); break; + case 19: launchFusedRouterMma<19>(params, data.mUsePdl, s); break; + case 20: launchFusedRouterMma<20>(params, data.mUsePdl, s); break; + case 21: launchFusedRouterMma<21>(params, data.mUsePdl, s); break; + case 22: launchFusedRouterMma<22>(params, data.mUsePdl, s); break; + case 23: launchFusedRouterMma<23>(params, data.mUsePdl, s); break; + default: launchFusedRouterMma<24>(params, data.mUsePdl, s); break; + } +} + +} // namespace +} // namespace routingCustom +} // namespace moe::dev::routing + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +void invokeFusedRouter(void const* routerHiddenStates, void const* routerWeight, int32_t hiddenDim, + int32_t* routerWorkspace, void* routingBias, batchedGemm::trtllm::gen::Dtype dtypeRoutingBias, int32_t numTokens, + int32_t numExperts, int32_t topK, int32_t localExpertOffset, int32_t localNumExperts, float routedScalingFactor, + int32_t tileTokensDim, int32_t* permutedIdxSize, int32_t* expandedIdxToPermutedIdx, int32_t* permutedIdxToTokenIdx, + void* expertWeights, int32_t* ctaIdxXyToBatchIdx, int32_t* ctaIdxXyToMnLimit, int32_t* numNonExitingCtas, + cudaStream_t stream) +{ + int32_t paddingLog2 = 0; + while ((1 << paddingLog2) < tileTokensDim) + { + ++paddingLog2; + } + TLLM_CHECK_WITH_INFO( + (1 << paddingLog2) == tileTokensDim, "tileTokensDim must be a power of 2, got %d", tileTokensDim); + + // The routing configuration of the TRTLLM-Gen MoE runner's DeepSeek-V3 routing without groups. + ::moe::dev::routing::routingCustom::Data routingData; + routingData.mDtypeOutput = batchedGemm::trtllm::gen::Dtype::Bfloat16; + routingData.mDtypeInput = batchedGemm::trtllm::gen::Dtype::Fp32; + routingData.mUsePdl = tensorrt_llm::common::getEnvEnablePDL(); + routingData.mPreprocessType = ::moe::dev::routing::RoutingPreprocessType::SigmoidBias; + routingData.mPostprocessType = ::moe::dev::routing::RoutingPostprocessType::ScaledSumNormalize; + routingData.mPtrRoutingBias = routingBias; + routingData.mDtypeBias = dtypeRoutingBias; + routingData.mRouteScale = routedScalingFactor; + routingData.mPtrPermutedIdxSize = permutedIdxSize; + routingData.mPtrExpandedIdxToPermutedIdx = expandedIdxToPermutedIdx; + routingData.mPtrPermutedIdxToExpandedIdx = nullptr; + routingData.mPtrPermutedIdxToTokenIdx = permutedIdxToTokenIdx; + routingData.mPtrTopKWeights = expertWeights; + routingData.mPtrCtaIdxXyToBatchIdx = ctaIdxXyToBatchIdx; + routingData.mPtrCtaIdxXyToMnLimit = ctaIdxXyToMnLimit; + routingData.mPtrNumNonExitingCtas = numNonExitingCtas; + routingData.mNumTokens = numTokens; + routingData.mNumExperts = numExperts; + routingData.mTopK = topK; + routingData.mPaddingLog2 = paddingLog2; + routingData.mTileTokensDim = tileTokensDim; + routingData.mLocalExpertsStartIdx = localExpertOffset; + routingData.mLocalExpertsStrideLog2 = 0; + routingData.mNumLocalExperts = localNumExperts; + + ::moe::dev::routing::routingCustom::runFusedRouter( + routingData, routerHiddenStates, routerWeight, hiddenDim, routerWorkspace, stream); +} + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMoeFusedRouter.h b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeFusedRouter.h similarity index 81% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmMoeFusedRouter.h rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeFusedRouter.h index 90b731dd876c..88a617474fc3 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMoeFusedRouter.h +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeFusedRouter.h @@ -29,16 +29,19 @@ namespace kernels::glm_moe_dsa { //! Workspace of invokeFusedRouter, in 32-bit words: two exchange buffers of [8 tokens][256 experts][4 slices] fp32 -//! whose idle value is -0.0 (bits 0x80000000), then the epoch (0 before the first launch), padded to 16 bytes. +//! (wide calls reuse each buffer for [tokens][256] reduced logits), whose idle value is -0.0 (bits 0x80000000), then +//! the epoch (0 before the first launch), padded to 16 bytes. constexpr int32_t kFusedRouterWorkspaceWords = 2 * 8 * 256 * 4 + 4; -//! The router GEMV and GLM-5's routing in one kernel. The logits are those of dsv3_router_gemm_op (bf16 -//! routerHiddenStates [numTokens, 6144] x routerWeight [256, 6144], fp32 logits, which are not stored); the routing is +//! The router product and GLM-5's routing in one kernel (bf16 +//! routerHiddenStates [numTokens, 6144] x routerWeight [256, 6144], fp32 logits, which are not returned). Through16 +//! tokens the logits match the router GEMV;17-24 use the owned tensor-core reduction. The routing is //! the TRTLLM-Gen MoE runner's DeepSeek-V3 routing without groups (sigmoid + bias, top-8, scaled sum normalization), //! with the outputs of its small-token routing kernel: the bf16 expert weights and the permutation and token-tile //! tables the TRTLLM-Gen FC1 and FC2 kernels read. //! -//! Covers hiddenDim 6144, 256 experts, top-8 and 1-8 tokens, at most tileTokensDim; anything else fails a check. +//! Covers hiddenDim6144,256 experts,top-8:1-8 tokens require tileTokensDim >= numTokens;9-24 support +//! multiple tiles per expert. Anything else fails a check. //! routerWorkspace: kFusedRouterWorkspaceWords words, owned by the caller and initialized once (both exchange buffers //! -0.0, epoch 0); the kernel leaves it in that state for the next launch. Launches that share a workspace must run in //! stream order. diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMoeGemm2K128.cu b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeGemm2K128.cu similarity index 99% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmMoeGemm2K128.cu rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeGemm2K128.cu index 376cd7e49771..9caa0e9e184e 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMoeGemm2K128.cu +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeGemm2K128.cu @@ -38,7 +38,7 @@ // Warps: 0-3 epilogue (TMEM lanes 32 w .. 32 w + 31 = weight rows of each tile), 4 weight loads, 5 TMEM allocation, // activations and MMA, 6 weight-scale copies. -#include "tensorrt_llm/kernels/glmMoeDsa/glmMoeGemm2K128.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeGemm2K128.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaDriverWrapper.h" diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmMoeGemm2K128.h b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeGemm2K128.h similarity index 100% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmMoeGemm2K128.h rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeGemm2K128.h diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeGemm2K256.cu b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeGemm2K256.cu new file mode 100644 index 000000000000..876074a04d6c --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeGemm2K256.cu @@ -0,0 +1,708 @@ +/* + * 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. + */ + +// FC2 of the NVFP4 MoE at intermediate size 256 (K = 256) and tile N 8 (see glmMoeGemm2K256.h for the contract). +// +// Work. A unit is (token tile t, 128-row tile of the output features); the units are split evenly over a persistent +// grid of one CTA per SM, each CTA owning a contiguous run of units (one or two token tiles). The weight scales of the +// whole run (2 KB per unit, one bulk copy per token tile) and the weight tiles of its first kStages units (16 KB of +// E2M1 per unit, one SWIZZLE_128B TMA box) are requested before the grid-dependency wait; each later weight tile goes +// into the stage of the unit kStages before it once that unit's MMAs are complete. The activations of the run's token +// tiles (8 rows x 128 B) and their 8x4 block scales are loaded after the wait. +// +// MMA. tcgen05.mma kind::mxf4nvf4 block16, M = 128 (weight rows), N = 8 (token rows), four K = 64 steps per unit in K +// order, issued by the elected lane of the converged MMA warp, one commit per unit. The TRTLLM-Gen Gemm2 cubin +// (t128x8x256, m128x8x64) runs the same four steps on the same operands and scales, so the FP32 accumulators and the +// BF16 outputs are equal bitwise. The weight scales go into a ring of TMEM slots by tcgen05.cp (a separate warp, ahead +// of the MMAs), the token scales after the wait. +// +// Accumulators. Unit k of a token tile's run writes its 8 TMEM columns at base + stride * k, stride = the tile's valid +// rows rounded up to 2, 4 or 8 (accumulators are 2-column aligned): the valid rows of all units of the tile lie in +// stride * units consecutive columns, and the tail columns a unit shares with the next one are padding rows, written +// in issue order. The epilogue reads them in 16-column loads, each once its last unit is complete. +// +// Warps: 0-3 epilogue (TMEM lanes 32 w .. 32 w + 31 = weight rows of each tile), 4 weight loads, 5 TMEM allocation, +// activations and MMA, 6 weight-scale copies. + +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeGemm2K256.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaDriverWrapper.h" +#include "tensorrt_llm/common/cudaUtils.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::gemm2K256 +{ +namespace +{ + +constexpr int kTileRows = 128; +constexpr int kK = 256; +constexpr int kKSteps = kK / 64; +constexpr int kTokens = 8; +constexpr int kTileBytesA = kTileRows * kK / 2; // 16,384: 128 rows x 128 B, SWIZZLE_128B +constexpr int kTileBytesSf = kTileRows * kK / 16; // 2,048: four 512 B blocks of the 128x4 layout +constexpr int kActBytes = kTokens * kK / 2; // 1,024: one SWIZZLE_128B atom +constexpr int kSfbBytes = 2 * 512; // two 32 rows x 16 B tcgen05.cp sources per token tile +constexpr int kStages = 8; // weight tiles in flight per CTA +constexpr int kSfSlots = 14; // TMEM weight-scale slots of 16 columns +constexpr int kMaxUnits = 29; // units per CTA +constexpr int kMaxSlots = 2; // token tiles per CTA (a run of kMaxUnits units spans at most two) +constexpr int kMaxTokenTiles = 128; +constexpr int kThreads = 224; +constexpr uint32_t kTmemCols = 512; +constexpr uint32_t kTmemAcc = 0; // packed accumulators, at most 256 columns +constexpr uint32_t kTmemSfa = 256; // kSfSlots x 16 columns: per K step 4 columns of weight scales +constexpr uint32_t kTmemSfb = 480; // kMaxSlots x 8 columns: per K step a column of token scales and a zero column +// Instruction descriptor, kind::mxf4nvf4: A/B E2M1, UE4M3 scales, K-major, M = 128, N = 8. +constexpr uint32_t kIdesc = (1u << 7) | (1u << 10) | ((8u >> 3) << 17) | ((128u >> 4) << 24); +// K-major SWIZZLE_128B operand descriptor without its start address: 8-row x 128 B atoms 1,024 B apart, version 1. +constexpr uint64_t kDescSw128 = (1ull << 16) | (static_cast(1024 >> 4) << 32) | (1ull << 46) | (2ull << 61); +// tcgen05.cp 32x128b source descriptor without its start address: 8-row x 16 B core matrices 128 B apart. +constexpr uint64_t kDescSf = (static_cast(128 >> 4) << 32) | (1ull << 46); +constexpr int kSmemBytes + = kStages * kTileBytesA + kMaxUnits * kTileBytesSf + kMaxSlots * (kActBytes + kSfbBytes) + 1024; + +static_assert( + kTmemSfa >= kTmemAcc + 256 && kTmemSfb >= kTmemSfa + 16 * kSfSlots && kTmemSfb + 8 * kMaxSlots <= kTmemCols, + "TMEM regions overlap"); + +struct Params +{ + CUtensorMap tmW; // weights [experts * hidden rows][128 B], box {128, 128}, SWIZZLE_128B + uint8_t const* wSf; + uint8_t const* act; + uint8_t const* actSf; + float const* scale; + int32_t const* numNonExitingCtas; + int32_t const* ctaIdxXyToBatchIdx; + int32_t const* ctaIdxXyToMnLimit; + __nv_bfloat16* out; + int32_t hidden; + int32_t maxCtas; +}; + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1100 + +__device__ __forceinline__ uint32_t smemU32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ void mbarInit(uint64_t* bar, uint32_t count) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" ::"r"(smemU32(bar)), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarExpectTx(uint64_t* bar, uint32_t bytes) +{ + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(smemU32(bar)), "r"(bytes) : "memory"); +} + +__device__ __forceinline__ void mbarWait(uint64_t* bar, uint32_t parity) +{ + uint32_t ok = 0; + while (!ok) + { + asm volatile( + "{\n" + ".reg .pred p;\n" + "mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n" + "selp.u32 %0, 1, 0, p;\n" + "}\n" + : "=r"(ok) + : "r"(smemU32(bar)), "r"(parity) + : "memory"); + } +} + +__device__ __forceinline__ void tmaLoad2d(void* dst, void const* map, int x, int y, uint64_t* bar) +{ + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1, {%2, %3}], [%4];" ::"r"( + smemU32(dst)), + "l"(map), "r"(x), "r"(y), "r"(smemU32(bar)) + : "memory"); +} + +__device__ __forceinline__ void bulkLoad(void* dst, void const* src, uint32_t bytes, uint64_t* bar) +{ + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"(smemU32(dst)), + "l"(src), "r"(bytes), "r"(smemU32(bar)) + : "memory"); +} + +__device__ __forceinline__ int ldRelaxed(int32_t const* p) +{ + int v; + asm volatile("ld.relaxed.gpu.global.b32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ void namedBarSync(int id, int threads) +{ + asm volatile("bar.sync %0, %1;" ::"r"(id), "r"(threads) : "memory"); +} + +__device__ __forceinline__ void namedBarArrive(int id, int threads) +{ + asm volatile("bar.arrive %0, %1;" ::"r"(id), "r"(threads) : "memory"); +} + +__device__ __forceinline__ void fenceProxyAsyncSmem() +{ + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + +__device__ __forceinline__ void tcFenceBefore() +{ + asm volatile("tcgen05.fence::before_thread_sync;" ::: "memory"); +} + +__device__ __forceinline__ void tcFenceAfter() +{ + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); +} + +__device__ __forceinline__ void tmemAlloc(uint32_t* dst, uint32_t cols) +{ + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"(smemU32(dst)), "r"(cols) + : "memory"); +} + +__device__ __forceinline__ void tmemRelinquish() +{ + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;" ::: "memory"); +} + +__device__ __forceinline__ void tmemDealloc(uint32_t taddr, uint32_t cols) +{ + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" ::"r"(taddr), "r"(cols) : "memory"); +} + +// The following three are executed by a whole (converged) warp; its elected lane issues the instruction. +__device__ __forceinline__ void utccpElect(uint32_t taddr, uint64_t sdesc) +{ + asm volatile( + "{\n" + ".reg .pred e;\n" + "elect.sync _|e, 0xffffffff;\n" + "@e tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;\n" + "}\n" ::"r"(taddr), + "l"(sdesc) + : "memory"); +} + +__device__ __forceinline__ void mmaElect(uint32_t d, uint64_t a, uint64_t b, uint32_t sfa, uint32_t sfb, uint32_t acc) +{ + asm volatile( + "{\n" + ".reg .pred e, p;\n" + "elect.sync _|e, 0xffffffff;\n" + "setp.ne.b32 p, %6, 0;\n" + "@e tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%4], [%5], p;\n" + "}\n" ::"r"(d), + "l"(a), "l"(b), "r"(kIdesc), "r"(sfa), "r"(sfb), "r"(acc) + : "memory"); +} + +__device__ __forceinline__ void commitElect(uint64_t* bar) +{ + asm volatile( + "{\n" + ".reg .pred e;\n" + "elect.sync _|e, 0xffffffff;\n" + "@e tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];\n" + "}\n" ::"r"(smemU32(bar)) + : "memory"); +} + +// Columns [taddr, taddr + 16) of the warp's TMEM lane quarter. +__device__ __forceinline__ void tmemLd16(uint32_t taddr, float (&v)[16]) +{ + uint32_t r[16]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32 {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, " + "%15}, [%16];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), "=r"(r[8]), + "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]) + : "r"(taddr) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" + : "+r"(r[0]), "+r"(r[1]), "+r"(r[2]), "+r"(r[3]), "+r"(r[4]), "+r"(r[5]), "+r"(r[6]), "+r"(r[7]), + "+r"(r[8]), "+r"(r[9]), "+r"(r[10]), "+r"(r[11]), "+r"(r[12]), "+r"(r[13]), "+r"(r[14]), "+r"(r[15]) + : + : "memory"); +#pragma unroll + for (int i = 0; i < 16; ++i) + { + v[i] = __uint_as_float(r[i]); + } +} + +// Stores the valid rows of the units [kBase, kBase + 16 / kStride) of a token tile from 16 packed accumulator columns +// (column kStride * k + r = row r of unit kBase + k). rowPtr: row 0 of the tile at this thread's feature of the run's +// first unit. +template +__device__ __forceinline__ void storeChunk( + float const (&v)[16], __nv_bfloat16* rowPtr, size_t rowStride, int kBase, int runLen, int nv, float scale) +{ + constexpr int kUnits = 16 / kStride; +#pragma unroll + for (int k = 0; k < kUnits; ++k) + { + if (kBase + k < runLen) + { + __nv_bfloat16* o = rowPtr + (kBase + k) * kTileRows; +#pragma unroll + for (int r = 0; r < kStride; ++r) + { + if (r < nv) + { + o[r * rowStride] = __float2bfloat16_rn(v[kStride * k + r] * scale); + } + } + } + } +} + +__device__ __forceinline__ int unitStart(int c, int units, int grid) +{ + return static_cast(static_cast(c) * units / grid); +} + +// Value j of the four per-lane tile entries (tiles 32 j + lane), j uniform across the warp. +__device__ __forceinline__ int pick4(int const (&v)[4], int j) +{ + return j == 0 ? v[0] : (j == 1 ? v[1] : (j == 2 ? v[2] : v[3])); +} + +#endif + +__global__ void __launch_bounds__(kThreads, 1) gemm2K256Kernel(const __grid_constant__ Params p) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1100 + extern __shared__ __align__(1024) uint8_t smemRaw[]; + __shared__ __align__(8) uint64_t fullBar[kStages]; + __shared__ __align__(8) uint64_t sfaBar[kMaxSlots]; + __shared__ __align__(8) uint64_t sfaReady[kSfSlots]; + __shared__ __align__(8) uint64_t accFull[kMaxUnits]; + __shared__ uint32_t sTmem; + __shared__ int sRange[2]; + __shared__ int sExpert[kMaxSlots]; + __shared__ int sNv[kMaxSlots]; + __shared__ int sStride[kMaxSlots]; + __shared__ int sCol[kMaxSlots]; + + uint8_t* smem = reinterpret_cast((reinterpret_cast(smemRaw) + 1023) & ~uintptr_t(1023)); + uint8_t* smemA = smem; + uint8_t* smemSf = smemA + kStages * kTileBytesA; + uint8_t* smemAct = smemSf + kMaxUnits * kTileBytesSf; + uint8_t* smemSfb = smemAct + kMaxSlots * kActBytes; + + int const warp = threadIdx.x >> 5; + int const lane = threadIdx.x & 31; + int const tilesPerSlot = p.hidden / kTileRows; + + if (warp == 4 && lane == 0) + { + for (int i = 0; i < kStages; ++i) + { + mbarInit(&fullBar[i], 1); + } + for (int i = 0; i < kSfSlots; ++i) + { + mbarInit(&sfaReady[i], 1); + } + for (int i = 0; i < kMaxUnits; ++i) + { + mbarInit(&accFull[i], 1); + } + for (int j = 0; j < kMaxSlots; ++j) + { + mbarInit(&sfaBar[j], 1); + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + asm volatile("prefetch.tensormap [%0];" ::"l"(&p.tmW) : "memory"); + } + if (warp == 5) + { + tmemAlloc(&sTmem, kTmemCols); + tmemRelinquish(); + } + tcFenceBefore(); + __syncthreads(); + tcFenceAfter(); + uint32_t const tmem = sTmem; + + if (warp == 4) + { + // The routing outputs are complete when this kernel starts: its primary, the FC1 cubin, triggers its dependents + // only after its own grid-dependency wait on the routing kernel. One round trip for the tile count and the + // experts and valid rows of every token tile; the weights and scales are constants. + int const numTiles = ldRelaxed(p.numNonExitingCtas); + int ex[4]; + int mn[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + int const t = lane + 32 * j; + ex[j] = t < p.maxCtas ? ldRelaxed(p.ctaIdxXyToBatchIdx + t) : 0; + mn[j] = t < p.maxCtas ? ldRelaxed(p.ctaIdxXyToMnLimit + t) : 0; + } + int const units = numTiles * tilesPerSlot; + int const a = unitStart(blockIdx.x, units, gridDim.x); + int const b = unitStart(blockIdx.x + 1, units, gridDim.x); + // Publish the run (its experts, valid rows and accumulator columns) before issuing the loads. + if (b > a) + { + int const tA = a / tilesPerSlot; + int const tB = (b - 1) / tilesPerSlot; + int col = 0; + for (int t = tA; t <= tB; ++t) + { + int const e = __shfl_sync(0xffffffffu, pick4(ex, t >> 5), t & 31); + int const m = __shfl_sync(0xffffffffu, pick4(mn, t >> 5), t & 31); + int const nv = m - kTokens * t; + int const stride = nv <= 2 ? 2 : (nv <= 4 ? 4 : 8); + int const runLen = min(b, (t + 1) * tilesPerSlot) - max(a, t * tilesPerSlot); + if (lane == 0) + { + sExpert[t - tA] = e; + sNv[t - tA] = nv; + sStride[t - tA] = stride; + sCol[t - tA] = col; + } + // The next token tile starts at a 32-column boundary (aligned TMEM loads). + col = (col + stride * runLen + 31) & ~31; + } + } + if (lane == 0) + { + sRange[0] = a; + sRange[1] = b; + } + __syncwarp(); + namedBarArrive(1, kThreads); + if (lane == 0 && b > a) + { + // The block scales of every unit of the run, one bulk copy per token tile (first, for the scale-copy warp). + int const tA = a / tilesPerSlot; + int i = 0; + int slot = 0; + while (a + i < b) + { + int const u0 = a + i; + int const t = u0 / tilesPerSlot; + int const nt0 = u0 - t * tilesPerSlot; + int const n = min(b - u0, tilesPerSlot - nt0); + mbarExpectTx(&sfaBar[slot], static_cast(n * kTileBytesSf)); + bulkLoad(smemSf + i * kTileBytesSf, + p.wSf + (static_cast(sExpert[slot]) * tilesPerSlot + nt0) * kTileBytesSf, + static_cast(n * kTileBytesSf), &sfaBar[slot]); + i += n; + ++slot; + } + // The weight tiles: the first kStages now, each later one once the MMAs of the unit kStages before it, + // which read its stage, are complete. + for (i = 0; a + i < b; ++i) + { + int const s = i % kStages; + if (i >= kStages) + { + mbarWait(&accFull[i - kStages], 0); + } + int const u = a + i; + int const t = u / tilesPerSlot; + int const row = sExpert[t - tA] * p.hidden + (u - t * tilesPerSlot) * kTileRows; + mbarExpectTx(&fullBar[s], kTileBytesA); + tmaLoad2d(smemA + s * kTileBytesA, &p.tmW, 0, row, &fullBar[s]); + } + } + } + else + { + namedBarSync(1, kThreads); + } + int const a = sRange[0]; + int const b = sRange[1]; + int const n = b - a; + int const t0 = n > 0 ? a / tilesPerSlot : 0; + int const nt0 = a - t0 * tilesPerSlot; + int const numSlots = n > 0 ? (b - 1) / tilesPerSlot - t0 + 1 : 0; + + if (warp == 6) + { + // The weight scales of each unit into its TMEM slot once their bulk copy has landed and the MMAs of the unit + // kSfSlots before it, which read that slot, are complete (before the grid-dependency wait returns for the first + // kSfSlots units of a CTA that runs while the previous kernel does); the MMA warp waits on sfaReady[slot]. + int u = 0; + int nt = nt0; + for (int j = 0; j < numSlots; ++j) + { + int const runLen = min(n - u, tilesPerSlot - nt); + mbarWait(&sfaBar[j], 0); + tcFenceAfter(); + for (int k = 0; k < runLen; ++k) + { + int const i = u + k; + int const slot = i % kSfSlots; + if (i >= kSfSlots) + { + mbarWait(&accFull[i - kSfSlots], 0); + tcFenceAfter(); + } + uint32_t const src = smemU32(smemSf + i * kTileBytesSf); + uint32_t const dst = tmem + kTmemSfa + 16 * slot; +#pragma unroll + for (int kb = 0; kb < kKSteps; ++kb) + { + utccpElect(dst + 4 * kb, kDescSf | (((src + 512 * kb) & 0x3ffffu) >> 4)); + } + commitElect(&sfaReady[slot]); + } + u += runLen; + nt = 0; + } + } + + // Everything below reads FC1's output or writes this kernel's output (the weight-load and scale-copy warps read + // neither). The dependents may launch once every CTA is past its wait. + if (warp < 4 || warp == 5) + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } + if (threadIdx.x == 0) + { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } + + if (warp < 4) + { + // Epilogue, per token tile of the run, in 16-column chunks, each once its last unit is complete. + uint32_t const tq = tmem + (static_cast(32 * warp) << 16); + // The weight rows are shuffled for the 128-row epilogue tile: row 32 w + l of a tile holds its feature + // 32 w + 4 (l & 7) + l / 8. + int const feature = 32 * warp + (lane & 7) * 4 + (lane >> 3); + int runStart = 0; + for (int j = 0; j < numSlots; ++j) + { + int const t = t0 + j; + int const ntFirst = j == 0 ? nt0 : 0; + int const runLen = min(n - runStart, tilesPerSlot - ntFirst); + int const nv = sNv[j]; + int const stride = sStride[j]; + float const scale = __ldg(p.scale + sExpert[j]); + int const unitsPerChunk = 16 / stride; + __nv_bfloat16* const rowBase + = p.out + static_cast(t) * kTokens * p.hidden + ntFirst * kTileRows + feature; + for (int kBase = 0; kBase < runLen; kBase += unitsPerChunk) + { + mbarWait(&accFull[runStart + min(kBase + unitsPerChunk, runLen) - 1], 0); + tcFenceAfter(); + float v[16]; + tmemLd16(tq + kTmemAcc + sCol[j] + stride * kBase, v); + if (stride == 2) + { + storeChunk<2>(v, rowBase, p.hidden, kBase, runLen, nv, scale); + } + else if (stride == 4) + { + storeChunk<4>(v, rowBase, p.hidden, kBase, runLen, nv, scale); + } + else + { + storeChunk<8>(v, rowBase, p.hidden, kBase, runLen, nv, scale); + } + } + runStart += runLen; + } + } + else if (warp == 5 && n > 0) + { + // Activations of the run's token tiles (8 rows of K = 256 E2M1, 128 B each, SWIZZLE_128B) and their 8x4 block + // scales, all loaded in one round trip. The scales go through shared memory (rows 0-7 of two 32-row blocks: + // per K step a column of 4 scales and a zero column) into every TMEM lane quarter. + int const r8 = lane >> 2; + int const c2 = (lane & 3) * 2; + uint4 actV[kMaxSlots][2]; + uint32_t sf[kMaxSlots][kKSteps]; +#pragma unroll + for (int j = 0; j < kMaxSlots; ++j) + { + actV[j][0] = make_uint4(0, 0, 0, 0); + actV[j][1] = make_uint4(0, 0, 0, 0); +#pragma unroll + for (int kb = 0; kb < kKSteps; ++kb) + { + sf[j][kb] = 0; + } + if (j < numSlots) + { + uint8_t const* row = p.act + (static_cast(t0 + j) * kTokens + r8) * (kK / 2); + actV[j][0] = __ldcg(reinterpret_cast(row + c2 * 16)); + actV[j][1] = __ldcg(reinterpret_cast(row + (c2 + 1) * 16)); + if (lane < kTokens) + { +#pragma unroll + for (int kb = 0; kb < kKSteps; ++kb) + { + sf[j][kb] = __ldcg( + reinterpret_cast(p.actSf) + (kKSteps * (t0 + j) + kb) * kTokens + lane); + } + } + } + } +#pragma unroll + for (int j = 0; j < kMaxSlots; ++j) + { + if (j < numSlots) + { + uint8_t* const atom = smemAct + j * kActBytes + r8 * 128; + *reinterpret_cast(atom + ((c2 ^ r8) << 4)) = actV[j][0]; + *reinterpret_cast(atom + (((c2 + 1) ^ r8) << 4)) = actV[j][1]; + *reinterpret_cast(smemSfb + j * kSfbBytes + lane * 16) = make_uint4(sf[j][0], 0u, sf[j][1], 0u); + *reinterpret_cast(smemSfb + j * kSfbBytes + 512 + lane * 16) + = make_uint4(sf[j][2], 0u, sf[j][3], 0u); + } + } + fenceProxyAsyncSmem(); + __syncwarp(); + for (int j = 0; j < numSlots; ++j) + { + uint32_t const src = smemU32(smemSfb + j * kSfbBytes); + utccpElect(tmem + kTmemSfb + 8 * j, kDescSf | ((src & 0x3ffffu) >> 4)); + utccpElect(tmem + kTmemSfb + 8 * j + 4, kDescSf | (((src + 512) & 0x3ffffu) >> 4)); + } + + int slot = 0; + int nt = nt0; + uint32_t const aBase = smemU32(smemA); + uint32_t bBase = smemU32(smemAct); + uint32_t sfb = tmem + kTmemSfb; + uint32_t d = tmem + kTmemAcc + sCol[0]; + uint32_t stride = sStride[0]; + for (int i = 0; i < n; ++i) + { + if (nt == tilesPerSlot) + { + ++slot; + nt = 0; + bBase += kActBytes; + sfb += 8; + d = tmem + kTmemAcc + sCol[slot]; + stride = sStride[slot]; + } + int const s = i % kStages; + int const sfSlot = i % kSfSlots; + mbarWait(&fullBar[s], (i / kStages) & 1); + mbarWait(&sfaReady[sfSlot], (i / kSfSlots) & 1); + tcFenceAfter(); + uint32_t const sfa = tmem + kTmemSfa + 16 * sfSlot; + uint64_t const da = kDescSw128 | (((aBase + s * kTileBytesA) & 0x3ffffu) >> 4); + uint64_t const db = kDescSw128 | ((bBase & 0x3ffffu) >> 4); +#pragma unroll + for (int kb = 0; kb < kKSteps; ++kb) + { + // K step kb: 32 B further along the 128 B rows of both operands. + mmaElect(d, da + 2 * kb, db + 2 * kb, sfa + 4 * kb, sfb + 2 * kb, kb > 0 ? 1u : 0u); + } + // A commit per unit: it releases the unit's weight stage and scale slot and tells the epilogue. + commitElect(&accFull[i]); + ++nt; + d += stride; + } + } + + tcFenceBefore(); + __syncthreads(); + if (warp == 5) + { + tcFenceAfter(); + tmemDealloc(tmem, kTmemCols); + } +#endif +} + +void encodeWeightMap(CUtensorMap* map, void const* weights, uint64_t rows) +{ + cuuint64_t const dims[2] = {kK / 2, rows}; + cuuint64_t const strides[1] = {kK / 2}; + cuuint32_t const box[2] = {kK / 2, kTileRows}; + cuuint32_t const elementStrides[2] = {1, 1}; + TLLM_CU_CHECK(tensorrt_llm::common::CUDADriverWrapper::getInstance()->cuTensorMapEncodeTiled(map, + CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(weights), dims, strides, box, elementStrides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_L2_128B, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); +} + +} // namespace + +bool isSupported(int32_t hiddenSize, int32_t maxNumCtas) +{ + int const sm = tensorrt_llm::common::getSMVersion(); + if (sm < 100 || sm >= 110 || hiddenSize % kTileRows != 0 || hiddenSize / kTileRows < kMaxUnits || maxNumCtas < 1 + || maxNumCtas > kMaxTokenTiles) + { + return false; + } + int const grid = tensorrt_llm::common::getMultiProcessorCount(); + int const maxUnits = maxNumCtas * (hiddenSize / kTileRows); + return (maxUnits + grid - 1) / grid <= kMaxUnits; +} + +void run(Data const& data, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(isSupported(data.hiddenSize, data.maxNumCtas), "gemm2K256: unsupported shape."); + TLLM_CUDA_CHECK(cudaFuncSetAttribute(gemm2K256Kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes)); + Params p{}; + encodeWeightMap(&p.tmW, data.weights, static_cast(data.numLocalExperts) * data.hiddenSize); + p.wSf = static_cast(data.weightsScale); + p.act = static_cast(data.act); + p.actSf = static_cast(data.actScale); + p.scale = data.outputScale; + p.numNonExitingCtas = data.numNonExitingCtas; + p.ctaIdxXyToBatchIdx = data.ctaIdxXyToBatchIdx; + p.ctaIdxXyToMnLimit = data.ctaIdxXyToMnLimit; + p.out = static_cast<__nv_bfloat16*>(data.output); + p.hidden = data.hiddenSize; + p.maxCtas = data.maxNumCtas; + + cudaLaunchConfig_t config{}; + config.gridDim = dim3(tensorrt_llm::common::getMultiProcessorCount()); + config.blockDim = dim3(kThreads); + config.dynamicSmemBytes = kSmemBytes; + config.stream = stream; + cudaLaunchAttribute attributes[2]{}; + attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attributes[0].val.programmaticStreamSerializationAllowed = data.usePdl ? 1 : 0; + attributes[1].id = cudaLaunchAttributePreferredSharedMemoryCarveout; + attributes[1].val.sharedMemCarveout = static_cast(data.sharedMemCarveout); + config.attrs = attributes; + config.numAttrs = data.sharedMemCarveout >= 0 ? 2 : 1; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, gemm2K256Kernel, p)); +} + +} // namespace kernels::glm_moe_dsa::gemm2K256 + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeGemm2K256.h b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeGemm2K256.h new file mode 100644 index 000000000000..fb35390d1e39 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeGemm2K256.h @@ -0,0 +1,67 @@ +/* + * 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. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::gemm2K256 +{ + +// FC2 of the NVFP4 MoE at intermediate size 256 and tile N 8, with the inputs and output of the TRTLLM-Gen Gemm2 +// cubins: +// output[8 t + r][f] = bf16(outputScale[e] * sum_k weights[e][f][k] * act[8 t + r][k]), e = ctaIdxXyToBatchIdx[t], +// for the routing's token tiles t < numNonExitingCtas and their valid rows r < ctaIdxXyToMnLimit[t] - 8 t. +struct Data +{ + // [numLocalExperts * hiddenSize][128 B] E2M1, the Gemm2 weight layout (rows shuffled for the 128-row epilogue + // tile). + void const* weights; + // E4M3 block scales of the weights, 128x4 layout: [numLocalExperts][hiddenSize / 128][4][512 B]. + void const* weightsScale; + // FC1 output: [rows][128 B] E2M1 and its E4M3 block scales in the 8x4 layout ([rows / 8][4][8][4]). + void const* act; + void const* actScale; + float const* outputScale; // [numLocalExperts] + int32_t const* numNonExitingCtas; + int32_t const* ctaIdxXyToBatchIdx; + int32_t const* ctaIdxXyToMnLimit; + void* output; // [rows][hiddenSize] BF16 + int32_t hiddenSize; + int32_t numLocalExperts; + int32_t maxNumCtas; // entries of ctaIdxXyToBatchIdx and ctaIdxXyToMnLimit + bool usePdl; + // The launch's preferred shared-memory carveout in percent (cudaLaunchAttributePreferredSharedMemoryCarveout; + // cudaSharedmemCarveoutMaxShared = 100), or -1: no preference, the driver chooses. + int32_t sharedMemCarveout = -1; +}; + +// Whether run() covers the shape on this device: an SM100-family GPU, hiddenSize a multiple of 128 and at least +// 3,712, at most 128 token tiles, and at most 29 units (token tile x 128 output features) per SM. +bool isSupported(int32_t hiddenSize, int32_t maxNumCtas); + +void run(Data const& data, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa::gemm2K256 + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmRouterGemm.cu b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmRouterGemm.cu new file mode 100644 index 000000000000..d982e6025bf9 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmRouterGemm.cu @@ -0,0 +1,536 @@ +/* + * 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/glmMoeDsaRubin/glmRouterGemm.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ + +namespace +{ + +// Convert 8 bfloat16 values from a uint4 to float array - optimized conversion +template +__device__ __forceinline__ void bf16_uint4_to_float8(uint4 const& vec, float* dst) +{ + __nv_bfloat16* bf16_ptr = reinterpret_cast<__nv_bfloat16*>(const_cast(&vec)); + +#pragma unroll + for (int i = 0; i < VPT; i++) + { + dst[i] = __bfloat162float(bf16_ptr[i]); + } +} + +// mat_a: the activations [kNumTokens, kHiddenDim]; mat_b: the weight [kNumExperts, kHiddenDim]; out: the logits +// [kNumTokens, kNumExperts]. Before the grid-dependency wait each CTA prefetches its weight row into L2 (no kernel of +// the step writes the weight), so the weight loads after the wait hit L2; the dependent grid may launch right after the +// wait instead of after the logits are written, and must wait on its own grid dependency before reading them. +template +__global__ __launch_bounds__(128, 1) void router_gemm_kernel(float* out, T const* mat_a, T const* mat_b) +{ + // Each block handles one expert column + int const n_idx = blockIdx.x; + int const tid = threadIdx.x; + constexpr int kWarpSize = 32; + constexpr int kNumWarps = kBlockSize / kWarpSize; + // Constants for this kernel + constexpr int k_elems_per_k_iteration = VPT * kBlockSize; + constexpr int k_iterations = kHiddenDim / k_elems_per_k_iteration; // Total K iterations + + // Initialize accumulators for all M rows + float acc[kNumTokens] = {}; + + // Shared memory for warp-level reduction + __shared__ float sm_reduction[kNumTokens][kNumWarps]; // kNumWarps + + // The weight row of expert n_idx (a column of B in column-major order) + T const* b_col = mat_b + n_idx * kHiddenDim; + + // Pre-compute k_base values for each iteration to help compiler optimize + int k_bases[k_iterations]; +#pragma unroll + for (int ki = 0; ki < k_iterations; ki++) + { + k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT; + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + { + // The primary grid does not write the weights, so this CTA's column can stream into L2 while the primary + // runs; the weight loads below then hit L2. What is loaded and the accumulation order are unchanged. + constexpr int kLineBytes = 128; + constexpr int kColumnBytes = kHiddenDim * static_cast(sizeof(T)); + static_assert(kColumnBytes % kLineBytes == 0, "the weight column must span whole 128-byte lines"); + char const* b_col_bytes = reinterpret_cast(b_col); + for (int offset = tid * kLineBytes; offset < kColumnBytes; offset += kBlockSize * kLineBytes) + { + asm volatile("prefetch.global.L2 [%0];" ::"l"(b_col_bytes + offset) : "memory"); + } + } + cudaGridDependencySynchronize(); + // Every dependent waits on its own grid dependency before it reads the logits, so it may launch now and overlap + // its prologue with this GEMV. + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + // Process the GEMM in chunks + for (int ki = 0; ki < k_iterations; ki++) + { + int const k_base = k_bases[ki]; + + // Load B matrix values using vector load (8 bf16 values) + uint4 b_vec = *reinterpret_cast(b_col + k_base); + + // Convert B values to float + float b_float[VPT]; + bf16_uint4_to_float8(b_vec, b_float); + +// Process each token +#pragma unroll + for (int m_idx = 0; m_idx < kNumTokens; m_idx++) + { + // Load both rows of A matrix using vector loads + uint4 a_vec = *reinterpret_cast(mat_a + (m_idx * kHiddenDim) + k_base); + + // Convert A values to float + float a_float[VPT]; + bf16_uint4_to_float8(a_vec, a_float); + +// Process elements in this chunk +#pragma unroll + for (int k = 0; k < VPT; k++) + { + float a = a_float[k]; + float b = b_float[k]; + acc[m_idx] += a * b; + } + } + } + + // Perform warp-level reduction + int const warpSize = 32; + int const warpId = tid / warpSize; + int const laneId = tid % warpSize; + + // Register for warp-level reduction results + float warp_result[kNumTokens]; + +#pragma unroll + for (int m_idx = 0; m_idx < kNumTokens; m_idx++) + { + warp_result[m_idx] = acc[m_idx]; + } + +// Perform warp-level reduction using optimized butterfly pattern +#pragma unroll + for (int m = 0; m < kNumTokens; m++) + { + float sum = warp_result[m]; + + // Butterfly reduction pattern + sum += __shfl_xor_sync(0xffffffff, sum, 16); + sum += __shfl_xor_sync(0xffffffff, sum, 8); + sum += __shfl_xor_sync(0xffffffff, sum, 4); + sum += __shfl_xor_sync(0xffffffff, sum, 2); + sum += __shfl_xor_sync(0xffffffff, sum, 1); + + // Only the first thread in each warp stores to shared memory + if (laneId == 0) + { + sm_reduction[m][warpId] = sum; + } + } + + __syncthreads(); + + // Final reduction across warps (only first thread) + if (tid == 0) + { +#pragma unroll + for (int m = 0; m < kNumTokens; m++) + { + float final_sum = 0.0f; + +// Sum across the kNumWarps +#pragma unroll + for (int w = 0; w < kNumWarps; w++) + { + final_sum += sm_reduction[m][w]; + } + + // Write final result + out[m * kNumExperts + n_idx] = final_sum; + } + } +} + +template +void launchRouterGemm( + float* logits, __nv_bfloat16 const* hiddenStates, __nv_bfloat16 const* weight, cudaStream_t stream) +{ + using T = __nv_bfloat16; + constexpr int VPT = 16 / sizeof(T); + constexpr int kBlockSize = 128; + cudaLaunchConfig_t config; + config.gridDim = kRouterNumExperts; + config.blockDim = kBlockSize; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.numAttrs = 1; + config.attrs = attrs; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, + router_gemm_kernel, logits, hiddenStates, + weight)); +} + +// --------------------------------------------------------------------------------------------------------------------- +// Router GEMM on tensor cores for 17-32 tokens, 256 experts, hidden 6144. +// +// A CTA owns 16 experts and one of 8 K slices of 768; the 8 CTAs of a cluster own the same experts. The CTA copies its +// weight tile (16 x 768) and the slice of every token's activations to shared memory, and each of its 4 warps runs +// mma.sync m16n8k16 (bf16 inputs, fp32 accumulation) over its 12 k-steps. The 4 warp tiles are added in warp order, +// sent to the shared memory of the cluster's rank 0, and rank 0 adds the 8 slices in rank order and writes the logits. +// The summation order is fixed, so the logits are deterministic; it differs from cuBLAS and from router_gemm_kernel. +// The weight tile is copied before the grid-dependency wait (router_gemm_kernel prefetches its weights into L2 there), +// and the dependent grid may launch right after that wait. +namespace routerMma +{ +constexpr int kHiddenDim = 6144; +constexpr int kNumExperts = 256; +// Experts per CTA: the M of one mma tile. +constexpr int kExpertsPerCta = 16; +// K slices of 768, one per CTA of a cluster. +constexpr int kSlices = 8; +constexpr int kSliceK = kHiddenDim / kSlices; +// Row pitch of the shared-memory tiles: the 8 rows of an ldmatrix start in distinct banks. +constexpr int kRowElems = kSliceK + 8; +constexpr int kWarps = 4; +constexpr int kThreads = kWarps * 32; +constexpr int kKStepsPerWarp = kSliceK / 16 / kWarps; +// 16-byte chunks per row of a slice. +constexpr int kChunks = kSliceK / 8; +constexpr int kNumCtas = kNumExperts / kExpertsPerCta * kSlices; +constexpr int kMinTokens = 17; +constexpr int kMaxTokens = 32; + +__host__ __device__ constexpr int paddedTokens(int numTokens) +{ + return (numTokens + 7) / 8 * 8; +} + +// Weight tile and activation rows (bf16), the 4 warp tiles and the 8 slice tiles received by rank 0 (fp32). +__host__ __device__ constexpr int smemBytes(int numTokens) +{ + return (kExpertsPerCta + paddedTokens(numTokens)) * kRowElems * 2 + + (kWarps + kSlices) * kExpertsPerCta * paddedTokens(numTokens) * 4; +} +} // namespace routerMma + +namespace +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +__device__ __forceinline__ uint32_t sharedAddress(void const* ptr) +{ + return static_cast(__cvta_generic_to_shared(ptr)); +} + +__device__ __forceinline__ void copyAsync16(void* dst, void const* src) +{ + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ::"r"(sharedAddress(dst)), "l"(src) : "memory"); +} + +__device__ __forceinline__ void copyAsyncCommit() +{ + asm volatile("cp.async.commit_group;" ::: "memory"); +} + +__device__ __forceinline__ void copyAsyncWaitAll() +{ + asm volatile("cp.async.wait_group 0;" ::: "memory"); +} + +__device__ __forceinline__ void ldmatrixX4(uint32_t (&r)[4], void const* ptr) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(sharedAddress(ptr))); +} + +__device__ __forceinline__ void ldmatrixX2(uint32_t (&r)[2], void const* ptr) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];" + : "=r"(r[0]), "=r"(r[1]) + : "r"(sharedAddress(ptr))); +} + +__device__ __forceinline__ void mmaBf16(float (&c)[4], uint32_t const (&a)[4], uint32_t const (&b)[2]) +{ + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); +} + +__device__ __forceinline__ void clusterArriveRelaxed() +{ + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); +} + +__device__ __forceinline__ void clusterWait() +{ + asm volatile("barrier.cluster.wait.aligned;" ::: "memory"); +} + +__device__ __forceinline__ void clusterArriveRelease() +{ + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); +} + +__device__ __forceinline__ void clusterWaitAcquire() +{ + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); +} + +// Stores v at the same shared-memory offset as ptr in the CTA of the given rank of this cluster. +__device__ __forceinline__ void storeToClusterRank(float* ptr, uint32_t rank, float v) +{ + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(sharedAddress(ptr)), "r"(rank)); + asm volatile("st.shared::cluster.f32 [%0], %1;" ::"r"(remote), "f"(v) : "memory"); +} +#endif +} // namespace + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) +template +__global__ void __launch_bounds__(routerMma::kThreads, 1) + router_gemm_mma_kernel(float* out, __nv_bfloat16 const* mat_a, __nv_bfloat16 const* mat_b) +{ + using namespace routerMma; + constexpr int kPadded = paddedTokens(kNumTokens); + constexpr int kTokenTiles = kPadded / 8; + constexpr int kTile = kExpertsPerCta * kPadded; + // Shared memory: the weight tile [kExpertsPerCta][kRowElems] and the activation rows [kPadded][kRowElems] (bf16), + // the warp tiles [kWarps][kExpertsPerCta][kPadded] and, used in rank 0, the slice tiles + // [kSlices][kExpertsPerCta][kPadded] (fp32). + extern __shared__ __align__(128) uint8_t smemMma[]; + __nv_bfloat16* weightTile = reinterpret_cast<__nv_bfloat16*>(smemMma); + __nv_bfloat16* actTile = weightTile + kExpertsPerCta * kRowElems; + float* warpSums = reinterpret_cast(actTile + kPadded * kRowElems); + float* sliceSums = warpSums + kWarps * kTile; + + // First phase of the cluster barrier: arrive now, wait before the first remote store, so that every CTA of the + // cluster has started by then. + clusterArriveRelaxed(); + int const slice = blockIdx.x % kSlices; // rank in the cluster + int const expertBase = blockIdx.x / kSlices * kExpertsPerCta; + int const warp = threadIdx.x / 32; + int const lane = threadIdx.x % 32; + + auto const copyWeightTile = [&]() + { + for (int c = threadIdx.x; c < kExpertsPerCta * kChunks; c += kThreads) + { + int const row = c / kChunks; + int const chunk = c % kChunks; + copyAsync16(weightTile + row * kRowElems + chunk * 8, + mat_b + static_cast(expertBase + row) * kHiddenDim + slice * kSliceK + chunk * 8); + } + copyAsyncCommit(); + }; + if constexpr (kPrefetchWeights) + { + // The primary grid does not write the weights. + copyWeightTile(); + } + // Rows past the last token are zero; they only feed outputs that are not stored. + for (int c = threadIdx.x; c < (kPadded - kNumTokens) * kChunks; c += kThreads) + { + *reinterpret_cast(actTile + (kNumTokens + c / kChunks) * kRowElems + (c % kChunks) * 8) + = make_uint4(0, 0, 0, 0); + } + cudaGridDependencySynchronize(); + if constexpr (kEarlyTrigger) + { + // Every dependent waits on its own grid dependency before it reads the logits. + cudaTriggerProgrammaticLaunchCompletion(); + } + if constexpr (!kPrefetchWeights) + { + copyWeightTile(); + } + for (int c = threadIdx.x; c < kNumTokens * kChunks; c += kThreads) + { + int const row = c / kChunks; + int const chunk = c % kChunks; + copyAsync16(actTile + row * kRowElems + chunk * 8, + mat_a + static_cast(row) * kHiddenDim + slice * kSliceK + chunk * 8); + } + copyAsyncCommit(); + copyAsyncWaitAll(); + __syncthreads(); + + float acc[kTokenTiles][4] = {}; +#pragma unroll + for (int s = 0; s < kKStepsPerWarp; s++) + { + int const k0 = (warp * kKStepsPerWarp + s) * 16; + uint32_t a[4]; + ldmatrixX4(a, weightTile + (lane % 16) * kRowElems + k0 + (lane / 16) * 8); +#pragma unroll + for (int j = 0; j < kTokenTiles; j++) + { + uint32_t b[2]; + ldmatrixX2(b, actTile + (j * 8 + lane % 8) * kRowElems + k0 + ((lane / 8) % 2) * 8); + mmaBf16(acc[j], a, b); + } + } + // m16n8 accumulators: acc[j][0..1] at expert lane / 4, tokens 8 j + 2 (lane % 4) + 0..1; acc[j][2..3] 8 experts on. +#pragma unroll + for (int j = 0; j < kTokenTiles; j++) + { + float* row = warpSums + (warp * kExpertsPerCta + lane / 4) * kPadded + j * 8 + 2 * (lane % 4); + *reinterpret_cast(row) = make_float2(acc[j][0], acc[j][1]); + *reinterpret_cast(row + 8 * kPadded) = make_float2(acc[j][2], acc[j][3]); + } + __syncthreads(); + clusterWait(); + for (int i = threadIdx.x; i < kTile; i += kThreads) + { + if (i % kPadded < kNumTokens) + { + float const sum = ((warpSums[i] + warpSums[kTile + i]) + warpSums[2 * kTile + i]) + warpSums[3 * kTile + i]; + storeToClusterRank(&sliceSums[slice * kTile + i], 0, sum); + } + } + // Second phase: rank 0 waits for the slices; the other CTAs are done (nothing reads their shared memory). + clusterArriveRelease(); + if (slice != 0) + { + return; + } + clusterWaitAcquire(); + for (int p = threadIdx.x; p < kExpertsPerCta * kNumTokens; p += kThreads) + { + int const e = p / kNumTokens; + int const m = p % kNumTokens; + float sum = sliceSums[e * kPadded + m]; +#pragma unroll + for (int r = 1; r < kSlices; r++) + { + sum += sliceSums[r * kTile + e * kPadded + m]; + } + out[m * kNumExperts + expertBase + e] = sum; + } + if constexpr (!kEarlyTrigger) + { + cudaTriggerProgrammaticLaunchCompletion(); + } +} +#else +template +__global__ void __launch_bounds__(routerMma::kThreads, 1) + router_gemm_mma_kernel(float* /* out */, __nv_bfloat16 const* /* mat_a */, __nv_bfloat16 const* /* mat_b */) +{ + assert(false && "router_gemm_mma_kernel is only supported on SM100+ architectures"); +} +#endif + +template +void invokeRouterGemmMma(float* output, T const* mat_a, T const* mat_b, cudaStream_t stream) +{ + static_assert(std::is_same_v, "the tensor-core router GEMM takes bf16 inputs"); + static_assert(kNumExperts == routerMma::kNumExperts && kHiddenDim == routerMma::kHiddenDim, + "the tensor-core router GEMM covers 256 experts and hidden size 6144"); + static_assert(kNumTokens >= routerMma::kMinTokens && kNumTokens <= routerMma::kMaxTokens, + "the tensor-core router GEMM covers 17-32 tokens"); + constexpr int kSmemBytes = routerMma::smemBytes(kNumTokens); + void (*kernel)(float*, __nv_bfloat16 const*, __nv_bfloat16 const*) = router_gemm_mma_kernel; + TLLM_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes)); + cudaLaunchConfig_t config; + config.gridDim = routerMma::kNumCtas; + config.blockDim = routerMma::kThreads; + config.dynamicSmemBytes = kSmemBytes; + config.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + attrs[1].id = cudaLaunchAttributeClusterDimension; + attrs[1].val.clusterDim.x = routerMma::kSlices; + attrs[1].val.clusterDim.y = 1; + attrs[1].val.clusterDim.z = 1; + config.numAttrs = 2; + config.attrs = attrs; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, output, mat_a, mat_b)); +} + +// Launches the instance of numTokens, one of kNumTokens..kRouterMaxTokens. +template +void dispatchRouterGemm( + float* logits, __nv_bfloat16 const* hiddenStates, __nv_bfloat16 const* weight, int numTokens, cudaStream_t stream) +{ + if constexpr (kNumTokens <= kRouterMaxTokens) + { + if (numTokens == kNumTokens) + { + if constexpr (kNumTokens <= 16) + { + launchRouterGemm(logits, hiddenStates, weight, stream); + } + else + { + invokeRouterGemmMma<__nv_bfloat16, kNumTokens, kRouterNumExperts, kRouterHiddenSize>( + logits, hiddenStates, weight, stream); + } + } + else + { + dispatchRouterGemm(logits, hiddenStates, weight, numTokens, stream); + } + } + else + { + TLLM_THROW("The GLM router GEMV takes 1 to %d tokens, got %d", kRouterMaxTokens, numTokens); + } +} + +} // namespace + +void invokeRouterGemm( + float* logits, __nv_bfloat16 const* hiddenStates, __nv_bfloat16 const* weight, int numTokens, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO( + numTokens >= 1, "The GLM router GEMV takes 1 to %d tokens, got %d", kRouterMaxTokens, numTokens); + dispatchRouterGemm<1>(logits, hiddenStates, weight, numTokens, stream); +} + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.h b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmRouterGemm.h similarity index 83% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.h rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmRouterGemm.h index c103c7eeb7e7..032c9b6e46e3 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.h +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmRouterGemm.h @@ -29,17 +29,17 @@ namespace kernels::glm_moe_dsa //! The router GEMV's K: the GLM-5 hidden size. constexpr int kRouterHiddenSize = 6144; -//! Routed experts: one CTA computes one expert's logits. +//! Routed experts. constexpr int kRouterNumExperts = 256; //! Tokens per call at most: one kernel instance per token count. -constexpr int kRouterMaxTokens = 16; +constexpr int kRouterMaxTokens = 32; //! logits[t, e] = sum_k hiddenStates[t, k] * weight[e, k]: BF16 operands, FP32 accumulation and FP32 logits, all -//! dense row-major. hiddenStates is [numTokens, 6144] with numTokens in [1, 16], weight [256, 6144] and logits -//! [numTokens, 256]. +//! dense row-major. hiddenStates is [numTokens, 6144] with numTokens in [1, 32], weight [256, 6144] and logits +//! [numTokens, 256].1-16 tokens use one GEMV CTA per expert;17-32 use eight-CTA tensor-core clusters. //! //! Launched with the programmatic-dependent-launch attribute when TRTLLM_ENABLE_PDL is set. Before its grid-dependency -//! wait each CTA prefetches its weight row into L2; right after the wait, before any logit is written, it triggers its +//! wait each CTA prefetches its weight row (GEMV) or stages its weight tile (MMA); right after the wait it triggers its //! dependents. Every kernel that follows on the stream must therefore wait on its own grid dependency before it reads //! the logits or writes hiddenStates. void invokeRouterGemm( diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.cu b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmSparseMlaFlashDecode.cu similarity index 54% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.cu rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmSparseMlaFlashDecode.cu index af9aade1887d..ac6e6d3289dd 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.cu +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmSparseMlaFlashDecode.cu @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmSparseMlaFlashDecode.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" @@ -34,7 +34,6 @@ namespace // The MMAs run 8 head columns: the kHeads query heads in columns 0 .. kHeads - 1, zero Q rows in the others. constexpr int kMmaHeads = 8; -constexpr int kHeads = kGlmSparseMlaNumHeads; constexpr int kHeadDimQk = kGlmSparseMlaHeadDimQk; constexpr int kHeadDimV = kGlmSparseMlaHeadDimV; // One cluster of kSplits CTAs per query token; CTA s owns indices [s * kTokensPerCta, (s + 1) * kTokensPerCta). @@ -44,7 +43,6 @@ constexpr int kThreads = 256; constexpr int kWarps = kThreads / 32; constexpr int kTokensPerWarp = kTokensPerCta / kWarps; static_assert(kTokensPerWarp == 16, "a warp gathers one 16-token MMA tile, one token row per lane"); -static_assert(kHeads < kMmaHeads, "the query heads fill part of the MMA head columns"); // A token row: 576 E4M3 values. In shared memory the rows are 592 bytes apart, so the eight 16-byte rows of every // ldmatrix phase fall into distinct banks. constexpr int kRowBytes = kHeadDimQk; @@ -76,26 +74,76 @@ constexpr int kOffRecvStats = kOffRecvO + kRecvOFloats * 4; constexpr int kRecvStatsFloats = kSplits * kMmaHeads * 2; constexpr int kRecvBytes = (kRecvOFloats + kRecvStatsFloats) * 4; constexpr int kOffVb = kOffRecvStats + kRecvStatsFloats * 4; -constexpr int kVbBytes = kHeads * kVbHeadBytes; -// The CTA's merged output slice in BF16, [8 head columns][32 columns]. -constexpr int kOffSlice = kOffVb + kVbBytes; +// Four heads use dedicated merge buffers. Eight heads reuse P and KV storage only after every +// CTA has finished P.V, as ordered by the post-P.V statistics transfer. constexpr int kSliceBytes = kMmaHeads * kDvPerCta * 2; -// The row's merged output in BF16, [8 head columns][512]. -constexpr int kOffGatheredO = kOffSlice + kSliceBytes; constexpr int kGatheredOBytes = kMmaHeads * kHeadDimV * 2; -constexpr int kOffBars = kOffGatheredO + kGatheredOBytes; -// One gather barrier per warp, the Q barrier, the receive barrier, the v_b slice barrier and the gathered-output -// barrier. constexpr int kNumBars = kWarps + 4; -constexpr int kSmemBytes = kOffBars + kNumBars * 8; -static_assert(kOffQ % 16 == 0 && kOffP % 16 == 0 && kOffRecvO % 16 == 0 && kOffVb % 16 == 0 && kOffSlice % 16 == 0 - && kOffGatheredO % 16 == 0 && kOffBars % 8 == 0, - "alignment"); +template +struct FoldedLayout +{ + static_assert(kHeads == kGlmSparseMlaNumHeads || kHeads == kGlmSparseMlaNumHeadsTp8); + static constexpr bool kOwnMergeBuffers = kHeads < kMmaHeads; + static constexpr int kVbBytes = kHeads * kVbHeadBytes; + static constexpr int kOffSlice = kOwnMergeBuffers ? kOffVb + kVbBytes : kOffP; + static constexpr int kOffGatheredO = kOwnMergeBuffers ? kOffSlice + kSliceBytes : 0; + static constexpr int kOffBars = kOwnMergeBuffers ? kOffGatheredO + kGatheredOBytes : kOffVb + kVbBytes; + static constexpr int kBytes = kOffBars + kNumBars * 8; + static_assert(kOffSlice % 16 == 0 && kOffGatheredO % 16 == 0 && kOffBars % 8 == 0, "alignment"); + static_assert(kBytes <= 227 * 1024, "shared memory per CTA"); +}; + +static_assert(kOffQ % 16 == 0 && kOffP % 16 == 0 && kOffRecvO % 16 == 0 && kOffVb % 16 == 0, "alignment"); static_assert(kSplits * kMmaHeads <= kThreads, "one max / sum transfer per thread"); static_assert(kMmaHeads * kDvPerCta == kThreads, "the merge maps one output element to each thread"); static_assert(kDvPerCta == 32 && kSplits <= 32, "the merge maps warp h to head h of the slice, lane r to partial r"); -static_assert(kSmemBytes <= 227 * 1024, "shared memory per CTA"); +static_assert(kGatheredOBytes <= kOffQ && kSliceBytes <= kMmaHeads * kPRowBytes, "merge buffers fit reused storage"); +static_assert(FoldedLayout<8>::kBytes == 230752, "eight-head shared-memory layout"); + +// The 16-head kernel (glmSparseMlaFlashDecodeHeads16Kernel): one cluster of kSplits CTAs per query token, CTA s owns +// indices [s * kTokensPerCta, (s + 1) * kTokensPerCta); each of its warps gathers one 16-token tile, as above. +// 16 splits (128 keys per CTA) while every cluster fits one wave; 8 splits (256 keys per CTA) for more query tokens. +template +struct Heads16 +{ + static constexpr int kHeads = kGlmSparseMlaNumHeadsWide; + // 8-head MMA tiles (the n dimension of every QK and P.V MMA). + static constexpr int kHeadTiles = kHeads / kMmaHeads; + static constexpr int kSplits = kSplits_; + static constexpr int kTokensPerCta = kGlmSparseMlaMaxTopK / kSplits; + static constexpr int kWarps = kTokensPerCta / kTokensPerWarp; + static constexpr int kThreads = kWarps * 32; + static constexpr int kPRowBytes = kTokensPerCta + 16; + // Output columns merged by each CTA of the cluster, and the 16-column P.V blocks of each warp. + static constexpr int kDvPerCta = kHeadDimV / kSplits; + static constexpr int kBlocksPerWarp = kHeadDimV / 16 / kWarps; + // The merge: thread t merges columns 2 (t % kColPairs) and 2 (t % kColPairs) + 1 of head t / kColPairs. + static constexpr int kColPairs = kDvPerCta / 2; + + // Shared memory: token rows, Q rows, P, the per-warp max / sum, the receive buffers of the cluster merge, + // mbarriers. + static constexpr int kOffQ = kTokensPerCta * kSmemRowBytes; + static constexpr int kOffP = kOffQ + kHeads * kSmemRowBytes; + static constexpr int kOffRed = kOffP + kHeads * kPRowBytes; + static constexpr int kRedFloats = 2 * kWarps * kHeads; + static constexpr int kOffRecvO = kOffRed + kRedFloats * 4; + static constexpr int kRecvOFloats = kSplits * kHeads * kDvPerCta; + static constexpr int kOffRecvStats = kOffRecvO + kRecvOFloats * 4; + static constexpr int kRecvStatsFloats = kSplits * kHeads * 2; + static constexpr int kRecvBytes = (kRecvOFloats + kRecvStatsFloats) * 4; + static constexpr int kOffBars = kOffRecvStats + kRecvStatsFloats * 4; + // One gather barrier per warp, the Q barrier and the receive barrier. + static constexpr int kNumBars = kWarps + 2; + static constexpr int kBytes = kOffBars + kNumBars * 8; + + static_assert(kTokensPerCta % kTokensPerWarp == 0 && kHeadDimV % (16 * kWarps) == 0, "whole MMA tiles per warp"); + static_assert(kDvPerCta % 16 == 0, "each output slice holds whole 16-column MMA blocks"); + static_assert(kOffQ % 16 == 0 && kOffP % 16 == 0 && kOffRecvO % 16 == 0 && kOffBars % 8 == 0, "alignment"); + static_assert(kSplits * kHeads <= kThreads, "one max / sum transfer per thread"); + static_assert(kHeads * kDvPerCta == 2 * kThreads, "the merge maps two output columns to each thread"); + static_assert(kBytes <= 227 * 1024, "shared memory per CTA"); +}; #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) @@ -252,6 +300,27 @@ __device__ __forceinline__ void mmaE4m3AsF16( bh[3]); } +// The A fragment of an E4M3 m16n8k32 MMA unpacked for two FP16 MMAs: ah[0..3] the low 16-bit halves of a0..a3, +// ah[4..7] the high halves. +__device__ __forceinline__ void unpackE4m3A(uint32_t (&ah)[8], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3) +{ + ah[0] = e4m3x2ToHalf2(a0); + ah[1] = e4m3x2ToHalf2(a1); + ah[2] = e4m3x2ToHalf2(a2); + ah[3] = e4m3x2ToHalf2(a3); + ah[4] = e4m3x2ToHalf2(a0 >> 16); + ah[5] = e4m3x2ToHalf2(a1 >> 16); + ah[6] = e4m3x2ToHalf2(a2 >> 16); + ah[7] = e4m3x2ToHalf2(a3 >> 16); +} + +// mmaE4m3AsF16 with A unpacked beforehand too (unpackE4m3A), so that one A fragment feeds several B fragments. +__device__ __forceinline__ void mmaE4m3AsF16Unpacked(float (&d)[4], uint32_t const (&ah)[8], uint32_t const (&bh)[4]) +{ + mmaF16(d, ah[0], ah[1], ah[2], ah[3], bh[0], bh[1]); + mmaF16(d, ah[4], ah[5], ah[6], ah[7], bh[2], bh[3]); +} + __device__ __forceinline__ uint8_t toE4m3(float x) { return static_cast(__nv_cvt_float_to_fp8(x, __NV_SATFINITE, __NV_E4M3)); @@ -321,11 +390,13 @@ __device__ __forceinline__ EarlyGatherRow gatherRowBeforeWait( // kRemap: the indices are positions, mapped to pool rows through the block table. kGatherBeforeWait (with kRemap): the // indices, the block table and the KV length are read, and the rows of the positions before this step's tokens // gathered, before the grid-dependency wait; the rows of this step's tokens after it. -template +template __global__ void __launch_bounds__(kThreads, 1) glmSparseMlaFlashDecodeKernel(GlmSparseMlaFlashDecodeParams const params) { #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) static_assert(kRemap || !kGatherBeforeWait, "the early gather tells this step's tokens by their positions"); + using Layout = FoldedLayout; + constexpr bool kStatsAfterPv = !Layout::kOwnMergeBuffers; extern __shared__ __align__(1024) uint8_t smem[]; uint8_t* sKv = smem; uint8_t* sQ = smem + kOffQ; @@ -334,9 +405,9 @@ __global__ void __launch_bounds__(kThreads, 1) glmSparseMlaFlashDecodeKernel(Glm float* sRedSum = sRedMax + kWarps * kMmaHeads; float* recvO = reinterpret_cast(smem + kOffRecvO); float* recvStats = reinterpret_cast(smem + kOffRecvStats); - uint8_t* sSlice = smem + kOffSlice; - uint8_t* sGatheredO = smem + kOffGatheredO; - uint64_t* bars = reinterpret_cast(smem + kOffBars); + uint8_t* sSlice = smem + Layout::kOffSlice; + uint8_t* sGatheredO = smem + Layout::kOffGatheredO; + uint64_t* bars = reinterpret_cast(smem + Layout::kOffBars); int const tid = threadIdx.x; int const warp = tid / 32; @@ -532,7 +603,7 @@ __global__ void __launch_bounds__(kThreads, 1) glmSparseMlaFlashDecodeKernel(Glm // issued once every token row has landed and arrives during QK / P.V. if (lane == 0) { - mbarrierArriveExpectTx(barVb, kVbBytes); + mbarrierArriveExpectTx(barVb, Layout::kVbBytes); } __syncwarp(); if (lane < kHeads) @@ -583,24 +654,26 @@ __global__ void __launch_bounds__(kThreads, 1) glmSparseMlaFlashDecodeKernel(Glm } __syncthreads(); - // This CTA's max and sum per head column go to every CTA of the cluster now; the partial output follows during - // P.V. The merged slice and the gathered output have buffers of their own, so no merge can overwrite memory that a - // warp of the cluster still reads in P.V. + // Dedicated merge buffers permit early statistics. Reused P/KV storage must wait until + // all P.V reads in the cluster are complete before any CTA can merge into it. clusterWait(); - if (tid < kSplits * kMmaHeads) + if constexpr (!kStatsAfterPv) { - int const dst = tid / kMmaHeads; - int const h = tid % kMmaHeads; - float ctaMax = -INFINITY; - float ctaSum = 0.f; -#pragma unroll - for (int w = 0; w < kWarps; ++w) + if (tid < kSplits * kMmaHeads) { - ctaMax = fmaxf(ctaMax, sRedMax[w * kMmaHeads + h]); - ctaSum += sRedSum[w * kMmaHeads + h]; + int const dst = tid / kMmaHeads; + int const h = tid % kMmaHeads; + float ctaMax = -INFINITY; + float ctaSum = 0.f; +#pragma unroll + for (int w = 0; w < kWarps; ++w) + { + ctaMax = fmaxf(ctaMax, sRedMax[w * kMmaHeads + h]); + ctaSum += sRedSum[w * kMmaHeads + h]; + } + storeRemote(mapToRank(smemAddress(recvStats + (split * kMmaHeads + h) * 2), dst), ctaMax, ctaSum, + mapToRank(barRecv, dst)); } - storeRemote(mapToRank(smemAddress(recvStats + (split * kMmaHeads + h) * 2), dst), ctaMax, ctaSum, - mapToRank(barRecv, dst)); } // O^T = V^T P^T: warp w owns the 16-column output blocks 4w .. 4w + 3, one block at a time over all k steps, each @@ -644,6 +717,25 @@ __global__ void __launch_bounds__(kThreads, 1) glmSparseMlaFlashDecodeKernel(Glm ob[1], ob[3], remoteBar); } } + if constexpr (kStatsAfterPv) + { + __syncthreads(); + if (tid < kSplits * kMmaHeads) + { + int const dst = tid / kMmaHeads; + int const h = tid % kMmaHeads; + float ctaMax = -INFINITY; + float ctaSum = 0.f; +#pragma unroll + for (int w = 0; w < kWarps; ++w) + { + ctaMax = fmaxf(ctaMax, sRedMax[w * kMmaHeads + h]); + ctaSum += sRedSum[w * kMmaHeads + h]; + } + storeRemote(mapToRank(smemAddress(recvStats + (split * kMmaHeads + h) * 2), dst), ctaMax, ctaSum, + mapToRank(barRecv, dst)); + } + } // The dependents (the o_proj GEMM) wait for the whole grid before reading out; launching them now lets them stage // their weights during the merge and the v_b product. cudaTriggerProgrammaticLaunchCompletion(); @@ -745,13 +837,397 @@ __global__ void __launch_bounds__(kThreads, 1) glmSparseMlaFlashDecodeKernel(Glm #endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } -template +// 16 query heads per token (Heads16 above): the 8-head kernel's gather, softmax and merge with two 8-head MMA tiles on +// every K / V fragment. QK and P.V issue the FP16 MMAs that the SM100 family runs for the E4M3 MMA (mmaE4m3AsF16), +// with the K / V fragment unpacked once for both head tiles. kGatherBeforeWait: as in sparseMlaFlashDecodeKernel. +template +__global__ void __launch_bounds__(Heads16::kThreads, 1) + glmSparseMlaFlashDecodeHeads16Kernel(GlmSparseMlaFlashDecodeParams const params) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + static_assert(kRemap || !kGatherBeforeWait, "the early gather tells this step's tokens by their positions"); + using h16 = Heads16; + extern __shared__ __align__(1024) uint8_t smem[]; + uint8_t* sKv = smem; + uint8_t* sQ = smem + h16::kOffQ; + uint8_t* sP = smem + h16::kOffP; + float* sRedMax = reinterpret_cast(smem + h16::kOffRed); + float* sRedSum = sRedMax + h16::kWarps * h16::kHeads; + float* recvO = reinterpret_cast(smem + h16::kOffRecvO); + float* recvStats = reinterpret_cast(smem + h16::kOffRecvStats); + uint64_t* bars = reinterpret_cast(smem + h16::kOffBars); + + int const tid = threadIdx.x; + int const warp = tid / 32; + int const lane = tid % 32; + // The cluster spans gridDim.x, so the split index is the CTA's rank in the cluster. + int const split = blockIdx.x; + int const tokenQ = blockIdx.y; + + uint32_t const barGather = smemAddress(bars + warp); + uint32_t const barQ = smemAddress(bars + h16::kWarps); + uint32_t const barRecv = smemAddress(bars + h16::kWarps + 1); + if (tid == 0) + { + for (int i = 0; i < h16::kNumBars; ++i) + { + mbarrierInit(smemAddress(bars + i), 1); + } + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); + mbarrierArriveExpectTx(barRecv, h16::kRecvBytes); + } + __syncthreads(); + // Paired with the wait before the first remote store: no CTA writes into a receive buffer whose barrier other + // CTAs of the cluster have not initialized yet. + clusterArriveRelaxed(); + + int const tokBase = warp * kTokensPerWarp; + [[maybe_unused]] EarlyGatherRow early{}; + if constexpr (kGatherBeforeWait) + { + if (warp == 0) + { + // The block-table row of the token's sequence into L2 while the indices load: the remap below reads it + // next. The CTAs of a cluster prefetch interleaved 128-byte lines of the row. + char const* tableRow = reinterpret_cast( + params.remapBlockTable + (tokenQ / params.seqLenQ) * params.remapBlockTableStride); + int64_t const rowBytes = static_cast(params.remapMaxBlocks) * sizeof(int32_t); + for (int64_t off = (static_cast(split) * 32 + lane) * 128; off < rowBytes; + off += h16::kSplits * 32 * 128) + { + asm volatile("prefetch.global.L2 [%0];\n" ::"l"(tableRow + off)); + } + } + early = gatherRowBeforeWait(params, tokenQ, split * h16::kTokensPerCta + tokBase + lane, lane, + smemAddress(sKv + (tokBase + lane) * kSmemRowBytes), barGather); + } + + // Q, the indices, the KV pool rows and the scales come from earlier kernels. + cudaGridDependencySynchronize(); + if (params.triggerAfterWait != 0) + { + // The dependent reads this grid's output only after its own grid-dependency wait, which waits for this + // whole grid; its CTAs may land and stage their weights while this kernel runs. + cudaTriggerProgrammaticLaunchCompletion(); + } + + if (warp == 0) + { + if (lane == 0) + { + mbarrierArriveExpectTx(barQ, h16::kHeads * kRowBytes); + } + __syncwarp(); + if (lane < h16::kHeads) + { + bulkCopyGlobalToShared(smemAddress(sQ + lane * kSmemRowBytes), + static_cast(params.q) + (static_cast(tokenQ) * h16::kHeads + lane) * kRowBytes, + kRowBytes, barQ); + } + } + + // Lane l of warp w gathers token tokBase + l of the CTA, as in sparseMlaFlashDecodeKernel. + uint32_t validMask; + if constexpr (kGatherBeforeWait) + { + // The rows of this step's tokens, which the KV append wrote. + if (lane < kTokensPerWarp && early.fresh) + { + bulkCopyGlobalToShared(smemAddress(sKv + (tokBase + lane) * kSmemRowBytes), + static_cast(params.kvPool) + early.row * kRowBytes, kRowBytes, barGather); + } + validMask = __ballot_sync(0xffffffffU, early.valid); + } + else + { + int const slot = split * h16::kTokensPerCta + tokBase + lane; + int indexRaw = -1; + if (lane < kTokensPerWarp && slot < params.topK) + { + indexRaw = __ldg(params.indices + static_cast(tokenQ) * params.indicesStride + slot); + } + int const seqIdx = tokenQ / params.seqLenQ; + int const tokenInSeq = tokenQ - seqIdx * params.seqLenQ; + if constexpr (kRemap) + { + // The index is a position in the sequence: map it to its pool row through the block table. + int const block = indexRaw / params.remapTokensPerBlock; + int page = -1; + if (indexRaw >= 0 && block < params.remapMaxBlocks) + { + page = __ldg(params.remapBlockTable + seqIdx * params.remapBlockTableStride + block); + } + int const inBlock = indexRaw - block * params.remapTokensPerBlock; + indexRaw = page >= 0 ? page * params.remapStrideFactor + inBlock + params.remapLayerOffset : -1; + } + int const numValid = min(__ldg(params.seqLensKv + seqIdx) - (params.seqLenQ - 1 - tokenInSeq), params.topK); + bool const valid = lane < kTokensPerWarp && slot < numValid && indexRaw >= 0; + // Masked lanes still load a row (row 0 of the pool) so that every gather barrier expects the same bytes. + int64_t const row = valid ? indexRaw : 0; + if (lane == 0) + { + mbarrierArriveExpectTx(barGather, kTokensPerWarp * kRowBytes); + } + __syncwarp(); + if (lane < kTokensPerWarp) + { + bulkCopyGlobalToShared(smemAddress(sKv + (tokBase + lane) * kSmemRowBytes), + static_cast(params.kvPool) + row * kRowBytes, kRowBytes, barGather); + } + validMask = __ballot_sync(0xffffffffU, valid); + } + float const scaleSoftmaxLog2 = __ldg(params.scaleSoftmaxLog2); + float const outputScale = __ldg(params.outputScale); + + // Fragment coordinates of mma.sync: g = row group, c = column pair. + int const g = lane / 4; + int const c = lane % 4; + + // S^T = K Q^T for the warp's 16 tokens and the 16 heads (head tile t = heads 8 t .. 8 t + 7); two accumulators per + // tile over alternating k steps. + mbarrierWait(barQ, 0); + mbarrierWait(barGather, 0); + float s[h16::kHeadTiles][4]; + { + float acc[h16::kHeadTiles][2][4]; +#pragma unroll + for (int t = 0; t < h16::kHeadTiles; ++t) + { +#pragma unroll + for (int i = 0; i < 4; ++i) + { + acc[t][0][i] = 0.f; + acc[t][1][i] = 0.f; + } + } + uint32_t const qAddr = smemAddress(sQ + (lane % 8) * kSmemRowBytes + 16 * ((lane / 8) % 2)); + uint32_t const kAddr + = smemAddress(sKv + (tokBase + lane % 8 + 8 * ((lane / 8) % 2)) * kSmemRowBytes + 16 * (lane / 16)); +#pragma unroll + for (int kd = 0; kd < kHeadDimQk / 32; ++kd) + { + uint32_t a0, a1, a2, a3; + ldmatrixX4(kAddr + 32 * kd, a0, a1, a2, a3); + uint32_t ah[8]; + unpackE4m3A(ah, a0, a1, a2, a3); +#pragma unroll + for (int t = 0; t < h16::kHeadTiles; ++t) + { + uint32_t b0, b1; + ldmatrixX2(qAddr + t * kMmaHeads * kSmemRowBytes + 32 * kd, b0, b1); + uint32_t bh[4]; + unpackE4m3B(bh, b0, b1); + mmaE4m3AsF16Unpacked(acc[t][kd % 2], ah, bh); + } + } + bool const valid0 = (validMask >> g) & 1U; + bool const valid1 = (validMask >> (g + 8)) & 1U; +#pragma unroll + for (int t = 0; t < h16::kHeadTiles; ++t) + { + s[t][0] = valid0 ? (acc[t][0][0] + acc[t][1][0]) * scaleSoftmaxLog2 : -INFINITY; + s[t][1] = valid0 ? (acc[t][0][1] + acc[t][1][1]) * scaleSoftmaxLog2 : -INFINITY; + s[t][2] = valid1 ? (acc[t][0][2] + acc[t][1][2]) * scaleSoftmaxLog2 : -INFINITY; + s[t][3] = valid1 ? (acc[t][0][3] + acc[t][1][3]) * scaleSoftmaxLog2 : -INFINITY; + } + } + + // The CTA's max per head; this thread holds heads 8 t + 2 c and 8 t + 2 c + 1. +#pragma unroll + for (int t = 0; t < h16::kHeadTiles; ++t) + { + float max0 = fmaxf(s[t][0], s[t][2]); + float max1 = fmaxf(s[t][1], s[t][3]); +#pragma unroll + for (int offset = 4; offset < 32; offset *= 2) + { + max0 = fmaxf(max0, __shfl_xor_sync(0xffffffffU, max0, offset)); + max1 = fmaxf(max1, __shfl_xor_sync(0xffffffffU, max1, offset)); + } + if (g == 0) + { + sRedMax[warp * h16::kHeads + kMmaHeads * t + 2 * c] = max0; + sRedMax[warp * h16::kHeads + kMmaHeads * t + 2 * c + 1] = max1; + } + } + // Every warp observed its own gather barrier before its QK; this barrier orders all token rows before the P.V + // reads of the other warps. + __syncthreads(); + + // P in E4M3 at the permuted k slots, and the per-head sums of the fp32 P. + { + int const slot0 = permutedSlot(tokBase + g); + int const slot1 = permutedSlot(tokBase + g + 8); +#pragma unroll + for (int t = 0; t < h16::kHeadTiles; ++t) + { + int const head0 = kMmaHeads * t + 2 * c; + float ctaMax0 = -INFINITY; + float ctaMax1 = -INFINITY; +#pragma unroll + for (int w = 0; w < h16::kWarps; ++w) + { + ctaMax0 = fmaxf(ctaMax0, sRedMax[w * h16::kHeads + head0]); + ctaMax1 = fmaxf(ctaMax1, sRedMax[w * h16::kHeads + head0 + 1]); + } + // A CTA without valid tokens keeps P = 0. + float const offset0 = ctaMax0 == -INFINITY ? 0.f : ctaMax0; + float const offset1 = ctaMax1 == -INFINITY ? 0.f : ctaMax1; + float const p00 = exp2f(s[t][0] - offset0); + float const p01 = exp2f(s[t][1] - offset1); + float const p10 = exp2f(s[t][2] - offset0); + float const p11 = exp2f(s[t][3] - offset1); + sP[head0 * h16::kPRowBytes + slot0] = toE4m3(p00 * kPScale); + sP[(head0 + 1) * h16::kPRowBytes + slot0] = toE4m3(p01 * kPScale); + sP[head0 * h16::kPRowBytes + slot1] = toE4m3(p10 * kPScale); + sP[(head0 + 1) * h16::kPRowBytes + slot1] = toE4m3(p11 * kPScale); + float sum0 = p00 + p10; + float sum1 = p01 + p11; +#pragma unroll + for (int offset = 4; offset < 32; offset *= 2) + { + sum0 += __shfl_xor_sync(0xffffffffU, sum0, offset); + sum1 += __shfl_xor_sync(0xffffffffU, sum1, offset); + } + if (g == 0) + { + sRedSum[warp * h16::kHeads + head0] = sum0; + sRedSum[warp * h16::kHeads + head0 + 1] = sum1; + } + } + } + __syncthreads(); + + // This CTA's max and sum per head go to every CTA of the cluster; the partial output follows after P.V. + clusterWait(); + if (tid < h16::kSplits * h16::kHeads) + { + int const dst = tid / h16::kHeads; + int const h = tid % h16::kHeads; + float ctaMax = -INFINITY; + float ctaSum = 0.f; +#pragma unroll + for (int w = 0; w < h16::kWarps; ++w) + { + ctaMax = fmaxf(ctaMax, sRedMax[w * h16::kHeads + h]); + ctaSum += sRedSum[w * h16::kHeads + h]; + } + storeRemote(mapToRank(smemAddress(recvStats + (split * h16::kHeads + h) * 2), dst), ctaMax, ctaSum, + mapToRank(barRecv, dst)); + } + + // O^T = V^T P^T: warp w owns the 16-column output blocks kBlocksPerWarp * w + bi, for both head tiles (see + // sparseMlaFlashDecodeKernel for the fragment layout). Each V fragment is unpacked once for the two tiles. + float o[h16::kBlocksPerWarp][h16::kHeadTiles][4]; +#pragma unroll + for (int bi = 0; bi < h16::kBlocksPerWarp; ++bi) + { +#pragma unroll + for (int t = 0; t < h16::kHeadTiles; ++t) + { +#pragma unroll + for (int i = 0; i < 4; ++i) + { + o[bi][t][i] = 0.f; + } + } + } + { + uint32_t const pAddr = smemAddress(sP + (lane % 8) * h16::kPRowBytes + 16 * ((lane / 8) % 2)); + uint32_t const vAddr = smemAddress(sKv + lane * kSmemRowBytes + 16 * (h16::kBlocksPerWarp * warp)); +#pragma unroll + for (int kStep = 0; kStep < h16::kTokensPerCta / 32; ++kStep) + { + uint32_t ph[h16::kHeadTiles][4]; +#pragma unroll + for (int t = 0; t < h16::kHeadTiles; ++t) + { + uint32_t pb0, pb1; + ldmatrixX2(pAddr + t * kMmaHeads * h16::kPRowBytes + 32 * kStep, pb0, pb1); + unpackE4m3B(ph[t], pb0, pb1); + } +#pragma unroll + for (int bi = 0; bi < h16::kBlocksPerWarp; ++bi) + { + uint32_t r0, r1, r2, r3; + ldmatrixX4Trans(vAddr + 32 * kStep * kSmemRowBytes + 16 * bi, r0, r1, r2, r3); + uint32_t ah[8]; + unpackE4m3A(ah, __byte_perm(r0, r1, 0x6420), __byte_perm(r0, r1, 0x7531), __byte_perm(r2, r3, 0x6420), + __byte_perm(r2, r3, 0x7531)); +#pragma unroll + for (int t = 0; t < h16::kHeadTiles; ++t) + { + mmaE4m3AsF16Unpacked(o[bi][t], ah, ph[t]); + } + } + } + } + + // Columns [kDvPerCta r, kDvPerCta (r + 1)) of the partial output go to CTA r. +#pragma unroll + for (int bi = 0; bi < h16::kBlocksPerWarp; ++bi) + { + int const col = 16 * (h16::kBlocksPerWarp * warp + bi) + 2 * g; + int const dst = col / h16::kDvPerCta; + int const colInSlice = col % h16::kDvPerCta; + uint32_t const remoteBar = mapToRank(barRecv, dst); +#pragma unroll + for (int t = 0; t < h16::kHeadTiles; ++t) + { + int const head0 = kMmaHeads * t + 2 * c; + storeRemote( + mapToRank(smemAddress(recvO + (split * h16::kHeads + head0) * h16::kDvPerCta + colInSlice), dst), + o[bi][t][0], o[bi][t][2], remoteBar); + storeRemote( + mapToRank(smemAddress(recvO + (split * h16::kHeads + head0 + 1) * h16::kDvPerCta + colInSlice), dst), + o[bi][t][1], o[bi][t][3], remoteBar); + } + } + + // Merge the kSplits partials of this CTA's output slice: thread t merges two columns of head t / kColPairs. + mbarrierWaitCluster(barRecv, 0); + { + int const h = tid / h16::kColPairs; + int const colInSlice = 2 * (tid % h16::kColPairs); + float mergedMax = -INFINITY; +#pragma unroll + for (int r = 0; r < h16::kSplits; ++r) + { + mergedMax = fmaxf(mergedMax, recvStats[(r * h16::kHeads + h) * 2]); + } + float mergedSum = 0.f; + float acc0 = 0.f; + float acc1 = 0.f; +#pragma unroll + for (int r = 0; r < h16::kSplits; ++r) + { + float const partMax = recvStats[(r * h16::kHeads + h) * 2]; + float const alpha = partMax == -INFINITY ? 0.f : exp2f(partMax - mergedMax); + mergedSum += alpha * recvStats[(r * h16::kHeads + h) * 2 + 1]; + float2 const part + = *reinterpret_cast(recvO + (r * h16::kHeads + h) * h16::kDvPerCta + colInSlice); + acc0 += alpha * part.x; + acc1 += alpha * part.y; + } + float const value0 = mergedSum > 0.f ? acc0 / (mergedSum * kPScale) * outputScale : 0.f; + float const value1 = mergedSum > 0.f ? acc1 / (mergedSum * kPScale) * outputScale : 0.f; + *reinterpret_cast<__nv_bfloat162*>(static_cast<__nv_bfloat16*>(params.out) + + static_cast(tokenQ) * params.outStride + h * kHeadDimV + split * h16::kDvPerCta + colInSlice) + = __floats2bfloat162_rn(value0, value1); + } + // The dependents may launch once every CTA has stored its slice. + cudaTriggerProgrammaticLaunchCompletion(); +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +template void launchGlmSparseMlaFlashDecode(GlmSparseMlaFlashDecodeParams const& params, cudaStream_t stream) { - auto const kernel = glmSparseMlaFlashDecodeKernel; + using Layout = FoldedLayout; + auto const kernel = glmSparseMlaFlashDecodeKernel; static bool const configured = [kernel]() { - TLLM_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes)); + TLLM_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, Layout::kBytes)); TLLM_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1)); return true; }(); @@ -760,7 +1236,7 @@ void launchGlmSparseMlaFlashDecode(GlmSparseMlaFlashDecodeParams const& params, cudaLaunchConfig_t config; config.gridDim = dim3(kSplits, params.numTokensQ, 1); config.blockDim = dim3(kThreads, 1, 1); - config.dynamicSmemBytes = kSmemBytes; + config.dynamicSmemBytes = Layout::kBytes; config.stream = stream; cudaLaunchAttribute attrs[2]; attrs[0].id = cudaLaunchAttributeClusterDimension; @@ -774,6 +1250,71 @@ void launchGlmSparseMlaFlashDecode(GlmSparseMlaFlashDecodeParams const& params, TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, params)); } +template +void launchGlmSparseMlaFlashDecodeHeads16(GlmSparseMlaFlashDecodeParams const& params, cudaStream_t stream) +{ + using Layout = Heads16; + auto const kernel = glmSparseMlaFlashDecodeHeads16Kernel; + static bool const configured = [kernel]() + { + TLLM_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, Layout::kBytes)); + TLLM_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1)); + return true; + }(); + (void) configured; + + cudaLaunchConfig_t config; + config.gridDim = dim3(Layout::kSplits, params.numTokensQ, 1); + config.blockDim = dim3(Layout::kThreads, 1, 1); + config.dynamicSmemBytes = Layout::kBytes; + config.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = Layout::kSplits; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.attrs = attrs; + config.numAttrs = 2; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, params)); +} + +template +void launchGlmSparseMlaFlashDecodeHeads16Splits( + GlmSparseMlaFlashDecodeParams const& params, bool remap, cudaStream_t stream) +{ + if (remap && params.gatherBeforeWait != 0) + { + launchGlmSparseMlaFlashDecodeHeads16(params, stream); + } + else if (remap) + { + launchGlmSparseMlaFlashDecodeHeads16(params, stream); + } + else + { + launchGlmSparseMlaFlashDecodeHeads16(params, stream); + } +} + +template +void launchGlmSparseMlaFlashDecodeFolded(GlmSparseMlaFlashDecodeParams const& params, bool remap, cudaStream_t stream) +{ + if (remap && params.gatherBeforeWait) + { + launchGlmSparseMlaFlashDecode(params, stream); + } + else if (remap) + { + launchGlmSparseMlaFlashDecode(params, stream); + } + else + { + launchGlmSparseMlaFlashDecode(params, stream); + } +} + } // namespace void invokeGlmSparseMlaFlashDecode(GlmSparseMlaFlashDecodeParams const& params, cudaStream_t stream) @@ -783,8 +1324,11 @@ void invokeGlmSparseMlaFlashDecode(GlmSparseMlaFlashDecodeParams const& params, TLLM_CHECK_WITH_INFO(params.seqLenQ > 0 && params.indicesStride >= params.topK && params.numTokensQ >= 0 && params.numTokensQ % params.seqLenQ == 0, "Invalid GLM sparse MLA geometry"); - TLLM_CHECK_WITH_INFO(params.vbProj != nullptr && params.out != nullptr && params.outStride >= kHeads * kVHeadDim, - "GLM sparse MLA flash decode needs v_b_proj and an output of %d columns per token", kHeads * kVHeadDim); + TLLM_CHECK_WITH_INFO(params.numHeadsQ == 4 || params.numHeadsQ == 8, "GLM folded sparse MLA requires 4 or 8 heads"); + TLLM_CHECK_WITH_INFO( + params.vbProj != nullptr && params.out != nullptr && params.outStride >= params.numHeadsQ * kVHeadDim, + "GLM sparse MLA flash decode needs v_b_proj and an output of %d columns per token", + params.numHeadsQ * kVHeadDim); if (params.numTokensQ == 0) { return; @@ -798,17 +1342,48 @@ void invokeGlmSparseMlaFlashDecode(GlmSparseMlaFlashDecodeParams const& params, } TLLM_CHECK_WITH_INFO(!params.gatherBeforeWait || remap, "The GLM sparse MLA flash decode gathers before its wait only with the index remap"); - if (remap && params.gatherBeforeWait) + if (params.numHeadsQ == kGlmSparseMlaNumHeads) { - launchGlmSparseMlaFlashDecode(params, stream); + launchGlmSparseMlaFlashDecodeFolded(params, remap, stream); } - else if (remap) + else + { + launchGlmSparseMlaFlashDecodeFolded(params, remap, stream); + } +} + +void invokeGlmSparseMlaFlashDecodeLatent(GlmSparseMlaFlashDecodeParams const& params, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(params.topK > 0 && params.topK <= kGlmSparseMlaMaxTopK, + "GLM sparse MLA flash decode supports TopK up to %d, got %d", kGlmSparseMlaMaxTopK, params.topK); + TLLM_CHECK_WITH_INFO(params.seqLenQ > 0 && params.indicesStride >= params.topK && params.numTokensQ >= 0 + && params.numTokensQ % params.seqLenQ == 0, + "Invalid GLM sparse MLA geometry"); + TLLM_CHECK_WITH_INFO(params.numTokensQ <= kGlmSparseMlaMaxWideTokens, + "GLM latent sparse MLA supports at most %d query tokens", kGlmSparseMlaMaxWideTokens); + TLLM_CHECK_WITH_INFO( + params.vbProj == nullptr && params.out != nullptr && params.outStride >= kGlmSparseMlaNumHeadsWide * kHeadDimV, + "GLM latent sparse MLA requires an unfused [tokens, 16, 512] output"); + if (params.numTokensQ == 0) + { + return; + } + bool const remap = params.remapBlockTable != nullptr; + if (remap) + { + TLLM_CHECK_WITH_INFO(params.remapTokensPerBlock > 0 && params.remapMaxBlocks > 0 + && params.remapBlockTableStride >= params.remapMaxBlocks, + "Invalid GLM sparse MLA index remap geometry"); + } + TLLM_CHECK_WITH_INFO(!params.gatherBeforeWait || remap, + "The GLM sparse MLA flash decode gathers before its wait only with the index remap"); + if (params.numTokensQ <= 8) { - launchGlmSparseMlaFlashDecode(params, stream); + launchGlmSparseMlaFlashDecodeHeads16Splits<16>(params, remap, stream); } else { - launchGlmSparseMlaFlashDecode(params, stream); + launchGlmSparseMlaFlashDecodeHeads16Splits<8>(params, remap, stream); } } diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.h b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmSparseMlaFlashDecode.h similarity index 72% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.h rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmSparseMlaFlashDecode.h index 7620932c477f..a8ce1088309b 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.h +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmSparseMlaFlashDecode.h @@ -27,10 +27,13 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels::glm_moe_dsa { -// The GLM-5.2 decode shape at TP16 the flash-decode sparse MLA serves: 4 query heads per rank sharing one latent KV -// head, QK head dim 576 (512 latent + 64 rope), latent V dim 512, a TopK of at most 2048 selected tokens per query +// The GLM-5.2 decode shape at TP16/TP8 the flash-decode sparse MLA serves: 4/8 query heads per rank sharing one latent +// KV head, QK head dim 576 (512 latent + 64 rope), latent V dim 512, a TopK of at most 2048 selected tokens per query // token, and the v_b BMM (v_head_dim 256) folded into the output. constexpr int32_t kGlmSparseMlaNumHeads = 4; +constexpr int32_t kGlmSparseMlaNumHeadsTp8 = 8; +constexpr int32_t kGlmSparseMlaNumHeadsWide = 16; +constexpr int32_t kGlmSparseMlaMaxWideTokens = 22; constexpr int32_t kGlmSparseMlaHeadDimQk = 576; constexpr int32_t kGlmSparseMlaHeadDimV = 512; constexpr int32_t kGlmSparseMlaMaxTopK = 2048; @@ -38,7 +41,7 @@ constexpr int32_t kGlmSparseMlaVHeadDim = 256; struct GlmSparseMlaFlashDecodeParams { - // [numTokensQ, 4, 576] E4M3. + // [numTokensQ, numHeadsQ, 576] E4M3 for folded output, [numTokensQ, 16, 576] for latent output. void const* q; // KV pool base: pool token row r holds 576 E4M3 values at byte offset r * 576. void const* kvPool; @@ -51,13 +54,14 @@ struct GlmSparseMlaFlashDecodeParams float const* scaleSoftmaxLog2; // Device scalar that scales the attention output. float const* outputScale; - // [4, 256, 512] BF16 v_b_proj. + // [numHeadsQ, 256, 512] BF16 v_b_proj for folded output; nullptr for latent output. void const* vbProj; // out[t][h * 256 + v] = sum_d O[t][h][d] * vbProj[h][v][d] (BF16, outStride elements between query tokens), where - // O is the attention output rounded to BF16. + // O is the attention output rounded to BF16. The latent entry writes O[t][16][512] directly. void* out; int64_t outStride; int32_t numTokensQ; + int32_t numHeadsQ = kGlmSparseMlaNumHeads; // Query tokens per sequence: query token t of a sequence attends to the first // min(seqLensKv - (seqLenQ - 1 - t), topK) of its indices. int32_t seqLenQ; @@ -80,26 +84,34 @@ struct GlmSparseMlaFlashDecodeParams // rows of the positions before this step's tokens (position < seqLensKv - seqLenQ, appended by earlier steps) // there too; the rows of this step's tokens, written by the KV append it depends on, after the wait. bool gatherBeforeWait; + // Latent16-head path only: release the dependent grid after our dependency wait, so a + // matched pre-wait v_b consumer may stage weights before waiting for this output. + bool triggerAfterWait = false; }; // Generation-only static-token sparse MLA with E4M3 Q and KV and the v_b BMM folded in, as flash decoding over the // gathered keys. Every query token gets one cluster of 16 CTAs, and CTA s owns indices [128 s, 128 s + 128): // - each warp gathers 16 token rows with one bulk copy per row onto its own mbarrier; // - S = Q K^T and P V run as FP16 mma.sync on E4M3 values unpacked exactly to FP16, with fp32 accumulation; the MMAs -// run 8 head columns, the 4 query heads in columns 0-3 and zero Q rows in 4-7, whose outputs are not stored; +// run 8 head columns; at TP16, columns 4-7 have zero Q rows and their outputs are not stored; // P = exp2(S * scale - max) is stored as E4M3 scaled by 448, as the TRTLLM-Gen generation kernels feed their second // GEMM, and the query fragments are unpacked once per warp; // - every CTA sends output slice r of its unnormalized partial output (one 16-column block at a time, during P.V), its // max and its sum straight into the shared memory of CTA r of the cluster (st.async), and CTA r merges the 16 // partials of its 32 output columns; // - CTA r also owns v in [16 r, 16 r + 16) of every head: it all-gathers the row's merged BF16 output (st.async), -// multiplies it with its v_b_proj slice (64 KiB, bulk-copied once every token row has landed) on mma.sync BF16, and -// writes those 64 values of out. -// Requires the SM100 family, numTokensQ * 16 CTAs of 170 KiB shared memory and clusters of 16 CTAs. PDL: the kernel -// reads global memory only after its grid-dependency wait (except the gatherBeforeWait reads above) and lets its +// multiplies it with its v_b_proj slice (64/128 KiB, copied after the token rows land) on mma.sync BF16, and +// writes 16 values per head. +// Requires the SM100 family and clusters of 16 CTAs, with 170 KiB (4 heads) or 226 KiB (8 heads) per CTA. PDL: the +// kernel reads global memory only after its grid-dependency wait (except the gatherBeforeWait reads above) and lets its // dependents launch once its partial outputs are sent. void invokeGlmSparseMlaFlashDecode(GlmSparseMlaFlashDecodeParams const& params, cudaStream_t stream); +// TP4 generation: 16 query heads, BF16 latent output [tokens, 16, 512], no v_b fold. +// At most22 query tokens;16 CTA splits for<=8 tokens, otherwise8 splits. The output +// has outStride elements between token rows and contiguous [16,512] head/latent axes. +void invokeGlmSparseMlaFlashDecodeLatent(GlmSparseMlaFlashDecodeParams const& params, cudaStream_t stream); + } // namespace kernels::glm_moe_dsa TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.cu b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmVerifyHeadAccept.cu similarity index 98% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.cu rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmVerifyHeadAccept.cu index 4fd734cab81d..93d996105032 100644 --- a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.cu +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmVerifyHeadAccept.cu @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmVerifyHeadAccept.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.h b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmVerifyHeadAccept.h similarity index 100% rename from cpp/tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.h rename to cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmVerifyHeadAccept.h diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmVerifyHeadSample.cu b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmVerifyHeadSample.cu new file mode 100644 index 000000000000..ba091182b93f --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmVerifyHeadSample.cu @@ -0,0 +1,1204 @@ +/* + * 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/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmVerifyHeadSample.h" +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ +namespace +{ +constexpr int kSampleBlockSize = 512; +constexpr int kSampleVecSize = 4; +constexpr int kSampleTileElts = kSampleBlockSize * kSampleVecSize; +constexpr int kSampleTilesPerChunk = 4; +// One CTA per chunk of a row in the multi-CTA passes. +constexpr int kSampleChunkElts = kSampleTileElts * kSampleTilesPerChunk; +constexpr int kPivotRadixBits = 8; +constexpr int kPivotRadix = 1 << kPivotRadixBits; +constexpr int kWarpSize = 32; +constexpr unsigned kFullWarpMask = 0xFFFFFFFFU; + +using SampleBlockReduce = cub::BlockReduce; +using SampleBlockScan = cub::BlockScan; + +int numSampleChunks(int vocabSize) +{ + return (vocabSize + kSampleChunkElts - 1) / kSampleChunkElts; +} + +// IEEE-754 pivot add/multiply without flush to zero. +__device__ __forceinline__ float ieeeAdd(float a, float b) +{ + float r; + asm("add.rn.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b)); + return r; +} + +__device__ __forceinline__ float ieeeMul(float a, float b) +{ + float r; + asm("mul.rn.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b)); + return r; +} + +// Order-preserving float-to-unsigned map for radix top-k. +__device__ __forceinline__ uint32_t toOrderedKey(float value) +{ + uint32_t const bits = __float_as_uint(value); + return (bits & 0x80000000U) ? ~bits : (bits ^ 0x80000000U); +} + +__device__ __forceinline__ float fromOrderedKey(uint32_t key) +{ + uint32_t const bits = (key & 0x80000000U) ? (key ^ 0x80000000U) : ~key; + return __uint_as_float(bits); +} + +__device__ __forceinline__ float inverseTemperature(float temperature) +{ + return temperature > 0.F ? 1.F / temperature : 1.F / 1.e-6F; +} + +// The value the softmax sees: the top-k mask keeps logits >= the k-th largest, then the logit is scaled by 1 / T. +__device__ __forceinline__ float maskedScaledLogit(float logit, bool masked, float pivot, float invTemp) +{ + float const kept = (masked && !(logit >= pivot)) ? -INFINITY : logit; + return __fmul_rn(kept, invTemp); +} + +struct FloatMax +{ + __device__ __forceinline__ float operator()(float a, float b) const + { + return fmaxf(a, b); + } +}; + +__device__ __forceinline__ float sampleWarpMax(float value) +{ +#pragma unroll + for (int mask = kWarpSize / 2; mask > 0; mask >>= 1) + { + value = fmaxf(value, __shfl_xor_sync(kFullWarpMask, value, mask)); + } + return value; +} + +__device__ __forceinline__ float sampleWarpSum(float value) +{ +#pragma unroll + for (int mask = kWarpSize / 2; mask > 0; mask >>= 1) + { + value += __shfl_xor_sync(kFullWarpMask, value, mask); + } + return value; +} + +// Counts this CTA in the row's counter once every thread's stores are done. Returns true in the last CTA of the row +// to arrive, which resets the counter and then sees the other CTAs' stores (read them with ld.cg). +__device__ __forceinline__ bool arriveLastOfRow(int32_t* rowCounter, int numChunks) +{ + __shared__ bool isLast; + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) + { + isLast = atomicAdd(rowCounter, 1) == numChunks - 1; + if (isLast) + { + *rowCounter = 0; + } + } + __syncthreads(); + if (isLast) + { + __threadfence(); + } + return isLast; +} + +// Direct sampling: the tokens of the kernels above drawn from per-segment softmax partials, without the probability +// tensor. A warp handles a segment of 512 logits (lane l holds elements 128 q + 4 l + j, q < 4, j < 4); a CTA of the +// first kernel handles a chunk of 16 segments. The second kernel draws on the first CTA of each row, whose threads +// hold up to kSegsPerThread segments each (1 up to 512 segments, else kDirectMaxSegsPerThread); with top_p it has one +// CTA per chunk, which only rows with the top-p test and without the top-k mask use. +constexpr int kDirectSegVecs = 4; +constexpr int kDirectSegElts = kWarpSize * kSampleVecSize * kDirectSegVecs; +constexpr int kDirectWarps = kSampleBlockSize / kWarpSize; +constexpr int kDirectMaxSegsPerThread = 4; +// Segments a warp loads before it reduces them in the per-row passes of the second kernel, and logit vectors a +// thread loads before it counts them in the top-k pivot passes. +constexpr int kDirectSegBatch = 4; +constexpr int kDirectPivotBatch = 8; +// Logits of a top-k row that the pivot search selects from instead of the whole row. +constexpr int kDirectPivotCandidates = 4096; +static_assert(kDirectSegElts == kSampleBlockSize, "the draw rescans a segment with one logit per thread"); +static_assert(kDirectSegElts * kDirectWarps == kSampleChunkElts, "a first-kernel CTA covers one sampling chunk"); +static_assert(kSampleBlockSize * kDirectMaxSegsPerThread * kDirectSegElts == (1 << 20), + "the draw holds at most kDirectMaxSegsPerThread segments per thread"); + +struct IntMax +{ + __device__ __forceinline__ int operator()(int a, int b) const + { + return max(a, b); + } +}; + +struct IntMin +{ + __device__ __forceinline__ int operator()(int a, int b) const + { + return min(a, b); + } +}; + +using SampleBlockReduceInt = cub::BlockReduce; +using DirectPivotBlockScan = cub::BlockScan; + +union DirectTempStorage +{ + SampleBlockReduce::TempStorage reduce; + SampleBlockReduceInt::TempStorage reduceInt; + SampleBlockScan::TempStorage scan; +}; + +struct DirectKernelParams +{ + float const* logits; + float const* temperatures; + int32_t const* topK; + float const* topP; + float const* minP; + bool perRowRng; + int64_t const* seed; + int64_t const* offset; + int32_t* rowCounters; // [numRows], zero before and after the call: the CTAs of a row's first top-p test + int32_t* tokens; + float* segMax; // [numRows, numSegs] max of v = masked logit / T over a segment + float* segSum; // [numRows, numSegs] sum of exp(v - segMax) over the segment + float* segAbove0; // [numRows, numSegs] sum of raw weights above the first pivot of the last top-p test + float* segAbove1; // [numRows, numSegs] sum of raw weights above its second pivot + float* laneMax; // [numRows, numSegs * 32] largest logit of each lane's 16 (rows with the top-k mask) + float* pivotPool; // [numRows, kDirectPivotCandidates] logits the top-k pivot is selected from + int vocabSize; + int numSegs; + int numChunks; // CTAs per row of the first kernel, 16 segments each +}; + +// The sampling parameters of a row. They are step inputs written before the forward pass, not by the kernel that +// precedes the sampler, so the direct-sampling kernels read them before the PDL wait. +struct DirectRow +{ + float invTemp; + float topP; + float minP; + float weightFloor; + int topK; + bool masked; // 0 < top_k < vocab: the top-k mask applies + bool nucleus; // top_p < 1: the top-p test applies (with top_p 1 the test passes on exact sums) +}; + +__device__ __forceinline__ DirectRow loadDirectRow(DirectKernelParams const& params, int row) +{ + DirectRow r; + r.invTemp = inverseTemperature(params.temperatures[row]); + r.topK = params.topK != nullptr ? params.topK[row] : 0; + r.topP = params.topP != nullptr ? params.topP[row] : 1.F; + r.masked = r.topK > 0 && r.topK < params.vocabSize; + r.nucleus = r.topP < 1.F; + r.minP = fminf(fmaxf(params.minP[row], 0.F), 1.F); + r.weightFloor = r.minP; + return r; +} + +__device__ __forceinline__ void initDirectRowRng( + curandStatePhilox4_32_10_t& state, DirectKernelParams const& params, int row) +{ + if (params.perRowRng) + { + curand_init(static_cast(params.seed[row]), + (1ULL << 63) | static_cast(params.offset[row]), 0, &state); + } + else + { + curand_init(static_cast(params.seed[0]), static_cast(row), + static_cast(params.offset[0]), &state); + } +} + +// All membership tests use raw weights relative to the row maximum. Normalizing individual +// weights can merge distinct FP32 weights into ties, changing FULL top-p/top-k support. +__device__ __forceinline__ float directWeight(float logit, DirectRow const& r, float rowMax) +{ + float const w = __expf(__fmul_rn(logit, r.invTemp) - rowMax); + return w >= r.weightFloor ? w : 0.F; +} + +// The 16 logits of segment `seg` that lane `lane` holds; valid[q] is false for the vectors past the row (or seg < 0). +__device__ __forceinline__ void loadSegment(float const* rowLogits, int vocabSize, int seg, int lane, + float4 (&x)[kDirectSegVecs], bool (&valid)[kDirectSegVecs]) +{ +#pragma unroll + for (int q = 0; q < kDirectSegVecs; ++q) + { + int const idx = seg * kDirectSegElts + (q * kWarpSize + lane) * kSampleVecSize; + valid[q] = seg >= 0 && idx < vocabSize; + x[q] = valid[q] ? *reinterpret_cast(rowLogits + idx) + : make_float4(-INFINITY, -INFINITY, -INFINITY, -INFINITY); + } +} + +// Max and sum of exp(v - max) of a loaded segment, v = masked logit / T (warp-wide); lane 0 stores them. A segment +// whose maximum is infinite adds nothing, as a chunk of the multi-CTA softmax. +__device__ __forceinline__ void storeSegmentPartials(DirectKernelParams const& params, int row, int seg, float invTemp, + bool masked, float pivot, float4 const (&x)[kDirectSegVecs], bool const (&valid)[kDirectSegVecs]) +{ + float values[kDirectSegVecs][kSampleVecSize]; + float maxValue = -INFINITY; +#pragma unroll + for (int q = 0; q < kDirectSegVecs; ++q) + { + float const xs[kSampleVecSize] = {x[q].x, x[q].y, x[q].z, x[q].w}; +#pragma unroll + for (int j = 0; j < kSampleVecSize; ++j) + { + values[q][j] = valid[q] ? maskedScaledLogit(xs[j], masked, pivot, invTemp) : -INFINITY; + maxValue = fmaxf(maxValue, values[q][j]); + } + } + maxValue = sampleWarpMax(maxValue); + bool const empty = isinf(maxValue); + float sum = 0.F; + if (!empty) + { +#pragma unroll + for (int q = 0; q < kDirectSegVecs; ++q) + { +#pragma unroll + for (int j = 0; j < kSampleVecSize; ++j) + { + sum += __expf(values[q][j] - maxValue); + } + } + } + sum = sampleWarpSum(sum); + if (threadIdx.x % kWarpSize == 0) + { + int64_t const idx = static_cast(row) * params.numSegs + seg; + params.segMax[idx] = empty ? -INFINITY : maxValue; + params.segSum[idx] = empty ? 0.F : sum; + } +} + +// The segments [segBegin, segEnd) of `row` that `take` selects, by the warps of one CTA: warp w tests segments +// segBegin + w, segBegin + w + 16, ... on their max of v (a lane each) and runs fn(seg, x, valid) on the selected +// ones, kDirectSegBatch loaded at a time; skip(seg) runs on the lane that tested a segment `take` rejected. +template +__device__ __forceinline__ void forEachSegment( + DirectKernelParams const& params, int row, int segBegin, int segEnd, Take take, Fn fn, Skip skip) +{ + int const lane = static_cast(threadIdx.x) % kWarpSize; + float const* rowLogits = params.logits + static_cast(row) * params.vocabSize; + int64_t const rowSegs = static_cast(row) * params.numSegs; + for (int base = segBegin + static_cast(threadIdx.x) / kWarpSize; base < segEnd; + base += kDirectWarps * kWarpSize) + { + int const tested = base + kDirectWarps * lane; + bool selected = false; + if (tested < segEnd) + { + selected = take(params.segMax[rowSegs + tested]); + if (!selected) + { + skip(tested); + } + } + unsigned pending = __ballot_sync(kFullWarpMask, selected); + while (pending != 0U) + { + int segs[kDirectSegBatch]; + float4 x[kDirectSegBatch][kDirectSegVecs]; + bool valid[kDirectSegBatch][kDirectSegVecs]; +#pragma unroll + for (int b = 0; b < kDirectSegBatch; ++b) + { + segs[b] = -1; + if (pending != 0U) + { + segs[b] = base + kDirectWarps * (__ffs(static_cast(pending)) - 1); + pending &= pending - 1U; + } + loadSegment(rowLogits, params.vocabSize, segs[b], lane, x[b], valid[b]); + } +#pragma unroll + for (int b = 0; b < kDirectSegBatch; ++b) + { + if (segs[b] >= 0) + { + fn(segs[b], x[b], valid[b]); + } + } + } + } +} + +// Save the original global maximum before filtered partials overwrite the unfiltered ones. +__device__ float originalRowMax(DirectKernelParams const& params, int row, DirectTempStorage& tempStorage) +{ + __shared__ float result; + float local = -INFINITY; + for (int seg = threadIdx.x; seg < params.numSegs; seg += kSampleBlockSize) + { + local = fmaxf(local, params.segMax[static_cast(row) * params.numSegs + seg]); + } + float const maximum = SampleBlockReduce(tempStorage.reduce).Reduce(local, FloatMax{}); + if (threadIdx.x == 0) + result = maximum; + __syncthreads(); + float const value = result; + __syncthreads(); + return value; +} + +// min-p and top-k compose in weight space before top-p. The original maximum always +// survives valid finite-input filters, including min_p=1 and ties at the kth weight. +__device__ void storeFilteredPartials(DirectKernelParams const& params, int row, DirectRow const& r, float rowMax) +{ + int64_t const rowSegs = static_cast(row) * params.numSegs; + forEachSegment( + params, row, 0, params.numSegs, [&](float maximum) { return __expf(maximum - rowMax) >= r.weightFloor; }, + [&](int seg, float4 const(&x)[kDirectSegVecs], bool const(&valid)[kDirectSegVecs]) + { + float values[kDirectSegVecs][kSampleVecSize]; + float maximum = -INFINITY; +#pragma unroll + for (int q = 0; q < kDirectSegVecs; ++q) + { + float const xs[kSampleVecSize] = {x[q].x, x[q].y, x[q].z, x[q].w}; +#pragma unroll + for (int j = 0; j < kSampleVecSize; ++j) + { + float const v = __fmul_rn(xs[j], r.invTemp); + bool const keep = valid[q] && __expf(v - rowMax) >= r.weightFloor; + values[q][j] = keep ? v : -INFINITY; + maximum = fmaxf(maximum, values[q][j]); + } + } + maximum = sampleWarpMax(maximum); + float sum = 0.F; + if (isfinite(maximum)) + { +#pragma unroll + for (int q = 0; q < kDirectSegVecs; ++q) + { +#pragma unroll + for (int j = 0; j < kSampleVecSize; ++j) + sum += __expf(values[q][j] - maximum); + } + } + sum = sampleWarpSum(sum); + if (threadIdx.x % kWarpSize == 0) + { + params.segMax[rowSegs + seg] = maximum; + params.segSum[rowSegs + seg] = sum; + } + }, + [&](int seg) + { + params.segMax[rowSegs + seg] = -INFINITY; + params.segSum[rowSegs + seg] = 0.F; + }); +} + +// A bounded-retry/rounding fallback. Logits may include masked -inf; a finite scaled maximum +// always return an argmax, which survives all supported filters. All nonfinite rows +// have no specified sampling distribution; return the first maximum when possible. +__device__ int rowArgmax(DirectKernelParams const& params, int row, DirectRow const& r) +{ + __shared__ SampleBlockReduce::TempStorage maxStorage; + __shared__ SampleBlockReduceInt::TempStorage minStorage; + __shared__ float maximum; + __shared__ int result; + float const* logits = params.logits + static_cast(row) * params.vocabSize; + float local = -INFINITY; + for (int i = threadIdx.x; i < params.vocabSize; i += kSampleBlockSize) + local = fmaxf(local, __fmul_rn(logits[i], r.invTemp)); + float const blockMax = SampleBlockReduce(maxStorage).Reduce(local, FloatMax{}); + if (threadIdx.x == 0) + maximum = blockMax; + __syncthreads(); + int candidate = INT_MAX; + for (int i = threadIdx.x; i < params.vocabSize; i += kSampleBlockSize) + if (__fmul_rn(logits[i], r.invTemp) == maximum) + candidate = min(candidate, i); + int const first = SampleBlockReduceInt(minStorage).Reduce(candidate, IntMin{}); + if (threadIdx.x == 0) + result = first == INT_MAX ? 0 : first; + __syncthreads(); + return result; +} + +// The k-th largest logit of a row (0 < k < vocabSize) by radix select over the ordered keys, on one CTA. Each thread +// loads kDirectPivotBatch vectors before it counts them. +__device__ float directTopKPivot(float const* rowLogits, int vocabSize, int k) +{ + __shared__ uint32_t histogram[kPivotRadix]; + __shared__ DirectPivotBlockScan::TempStorage scanStorage; + __shared__ uint32_t selectedDigit; + __shared__ uint32_t countAbove; + int const tid = static_cast(threadIdx.x); + auto const* vecs = reinterpret_cast(rowLogits); + int const numVecs = vocabSize / kSampleVecSize; + uint32_t prefix = 0U; + uint32_t prefixMask = 0U; + uint32_t remaining = static_cast(k); +#pragma unroll 1 + for (int shift = 32 - kPivotRadixBits; shift >= 0; shift -= kPivotRadixBits) + { + for (int i = tid; i < kPivotRadix; i += kSampleBlockSize) + { + histogram[i] = 0U; + } + __syncthreads(); + for (int base = tid; base < numVecs; base += kSampleBlockSize * kDirectPivotBatch) + { + float4 x[kDirectPivotBatch]; +#pragma unroll + for (int b = 0; b < kDirectPivotBatch; ++b) + { + int const v = base + b * kSampleBlockSize; + x[b] = v < numVecs ? vecs[v] : make_float4(0.F, 0.F, 0.F, 0.F); + } +#pragma unroll + for (int b = 0; b < kDirectPivotBatch; ++b) + { + if (base + b * kSampleBlockSize < numVecs) + { + float const elts[kSampleVecSize] = {x[b].x, x[b].y, x[b].z, x[b].w}; +#pragma unroll + for (int j = 0; j < kSampleVecSize; ++j) + { + uint32_t const key = toOrderedKey(elts[j]); + if ((key & prefixMask) == prefix) + { + atomicAdd(&histogram[(key >> shift) & (kPivotRadix - 1)], 1U); + } + } + } + } + } + __syncthreads(); + // Digits in descending order: the selected digit is the first whose running count reaches `remaining`. + uint32_t const count = tid < kPivotRadix ? histogram[kPivotRadix - 1 - tid] : 0U; + uint32_t inclusive; + DirectPivotBlockScan(scanStorage).InclusiveSum(count, inclusive); + if (tid < kPivotRadix && inclusive >= remaining && inclusive - count < remaining) + { + selectedDigit = kPivotRadix - 1 - tid; + countAbove = inclusive - count; + } + __syncthreads(); + prefix |= selectedDigit << shift; + prefixMask |= static_cast(kPivotRadix - 1) << shift; + remaining -= countAbove; + __syncthreads(); + } + return fromOrderedKey(prefix); +} + +// The k-th largest logit of a row with the top-k mask (0 < k < vocabSize), on one CTA. One pass stores the largest of +// each lane's 16 logits of a segment (radix order). T, the k-th largest of these lane maxima, is at most the pivot (k +// distinct logits reach it), and every logit >= T lies in a lane whose maximum is >= T: the pivot is the k-th largest +// of the logits >= T, gathered and selected from when there are at most kDirectPivotCandidates of them, else selected +// over the whole row. +__device__ float directPivot(DirectKernelParams const& params, int row, int k) +{ + __shared__ int numCandidates; + int const tid = static_cast(threadIdx.x); + int const lane = tid % kWarpSize; + float const* rowLogits = params.logits + static_cast(row) * params.vocabSize; + int const numLanes = params.numSegs * kWarpSize; + if (k <= numLanes) + { + float* candidates = params.pivotPool + static_cast(row) * kDirectPivotCandidates; + float* rowLaneMax = params.laneMax + static_cast(row) * numLanes; + forEachSegment( + params, row, 0, params.numSegs, [](float) { return true; }, + [&](int seg, float4 const(&x)[kDirectSegVecs], bool const(&valid)[kDirectSegVecs]) + { + // A lane past the row holds the smallest key. + uint32_t laneKey = 0U; +#pragma unroll + for (int q = 0; q < kDirectSegVecs; ++q) + { + if (valid[q]) + { + laneKey = max(laneKey, toOrderedKey(x[q].x)); + laneKey = max(laneKey, toOrderedKey(x[q].y)); + laneKey = max(laneKey, toOrderedKey(x[q].z)); + laneKey = max(laneKey, toOrderedKey(x[q].w)); + } + } + rowLaneMax[seg * kWarpSize + lane] = fromOrderedKey(laneKey); + }, + [](int) {}); + __syncthreads(); + uint32_t const boundKey = toOrderedKey(directTopKPivot(rowLaneMax, numLanes, k)); + if (tid == 0) + { + numCandidates = 0; + } + __syncthreads(); + for (int laneIdx = tid; laneIdx < numLanes; laneIdx += kSampleBlockSize) + { + if (toOrderedKey(rowLaneMax[laneIdx]) < boundKey) + { + continue; + } + int const laneSeg = laneIdx / kWarpSize; + int const laneInSeg = laneIdx % kWarpSize; +#pragma unroll + for (int q = 0; q < kDirectSegVecs; ++q) + { + int const idx = laneSeg * kDirectSegElts + (q * kWarpSize + laneInSeg) * kSampleVecSize; + if (idx < params.vocabSize) + { + float4 const x = *reinterpret_cast(rowLogits + idx); + float const xs[kSampleVecSize] = {x.x, x.y, x.z, x.w}; +#pragma unroll + for (int j = 0; j < kSampleVecSize; ++j) + { + if (toOrderedKey(xs[j]) >= boundKey) + { + int const slot = atomicAdd(&numCandidates, 1); + if (slot < kDirectPivotCandidates) + { + candidates[slot] = xs[j]; + } + } + } + } + } + } + __syncthreads(); + int const count = numCandidates; + if (count <= kDirectPivotCandidates) + { + // Pad to whole vectors with the smallest key: the k-th largest of count >= k candidates stays. + int const padded = (count + kSampleVecSize - 1) / kSampleVecSize * kSampleVecSize; + for (int i = count + tid; i < padded; i += kSampleBlockSize) + { + candidates[i] = fromOrderedKey(0U); + } + __syncthreads(); + return directTopKPivot(candidates, padded, k); + } + } + return directTopKPivot(rowLogits, params.vocabSize, k); +} + +// The last index of segments fromSeg, fromSeg - 1, ..., 0 whose exp(v - rowMax) is > 0 (searched backwards, one +// logit per thread) and that value; -1 when there is none. +__device__ int lastPositiveToken(DirectKernelParams const& params, int row, DirectRow const& r, float rowMax, + int fromSeg, float& tokenExp, DirectTempStorage& tempStorage) +{ + __shared__ int lastIdx; + __shared__ float lastExp; + float const* rowLogits = params.logits + static_cast(row) * params.vocabSize; + for (int seg = fromSeg; seg >= 0; --seg) + { + int const idx = seg * kDirectSegElts + static_cast(threadIdx.x); + float const e = idx < params.vocabSize ? directWeight(rowLogits[idx], r, rowMax) : 0.F; + int const blockLast = SampleBlockReduceInt(tempStorage.reduceInt).Reduce(e > 0.F ? idx : -1, IntMax{}); + if (threadIdx.x == 0) + { + lastIdx = blockLast; + } + __syncthreads(); + int const found = lastIdx; + if (found >= 0 && idx == found) + { + lastExp = e; + } + __syncthreads(); + if (found >= 0) + { + tokenExp = lastExp; + return found; + } + } + tokenExp = 0.F; + return -1; +} + +struct DirectDraw +{ + int token; // -1 when no probability is > 0 + float rowMax; // max of v over the row + float rowSum; // sum of exp(v - rowMax) over the row + float tokenExp; // exp(v - rowMax) of the token +}; + +// Inverse-CDF draw for `row` from its segment partials on one CTA: the first token +// whose running sum of exp(v - rowMax) exceeds u * rowSum, which is the first whose running probability exceeds u. +// The segment where the running sum passes u * rowSum comes from the partials, the token from that segment's 512 +// logits. u at or above the running sum gives the last positive-weight token. *u is in +// shared memory; it is read after the first barrier. draw.tokenExp is set when needTokenExp. +template +__device__ DirectDraw drawFromSegments(DirectKernelParams const& params, int row, DirectRow const& r, float const* u, + bool needTokenExp, DirectTempStorage& tempStorage) +{ + __shared__ float sharedMax; + __shared__ int crossSeg; + __shared__ float crossBefore; + __shared__ int sampled; + __shared__ float sampledExp; + int const tid = static_cast(threadIdx.x); + int const numSegs = params.numSegs; + int const perThread = (numSegs + kSampleBlockSize - 1) / kSampleBlockSize; + int64_t const rowSegs = static_cast(row) * numSegs; + // Thread t holds segments [t * perThread, (t + 1) * perThread). + float partMax[kSegsPerThread]; + float partSum[kSegsPerThread]; + float threadMax = -INFINITY; +#pragma unroll + for (int i = 0; i < kSegsPerThread; ++i) + { + int const seg = tid * perThread + i; + bool const valid = i < perThread && seg < numSegs; + partMax[i] = valid ? params.segMax[rowSegs + seg] : -INFINITY; + partSum[i] = valid ? params.segSum[rowSegs + seg] : 0.F; + threadMax = fmaxf(threadMax, partMax[i]); + } + if (tid == 0) + { + crossSeg = INT_MAX; + sampled = INT_MAX; + } + float const blockMax = SampleBlockReduce(tempStorage.reduce).Reduce(threadMax, FloatMax{}); + if (tid == 0) + { + sharedMax = blockMax; + } + __syncthreads(); + DirectDraw draw; + draw.rowMax = sharedMax; + draw.token = -1; + draw.tokenExp = 0.F; + float weight[kSegsPerThread]; + float threadSum = 0.F; +#pragma unroll + for (int i = 0; i < kSegsPerThread; ++i) + { + weight[i] = isinf(partMax[i]) ? 0.F : partSum[i] * __expf(partMax[i] - draw.rowMax); + threadSum += weight[i]; + } + float before; + SampleBlockScan(tempStorage.scan).ExclusiveSum(threadSum, before, draw.rowSum); + if (!(draw.rowSum > 0.F)) + { + // No probability is > 0: every logit masked or -inf, or NaN. + return draw; + } + float const target = *u * draw.rowSum; + float running = before; + int cross = INT_MAX; + float crossRunning = 0.F; +#pragma unroll + for (int i = 0; i < kSegsPerThread; ++i) + { + if (cross == INT_MAX && i < perThread && weight[i] > 0.F && running + weight[i] > target) + { + cross = tid * perThread + i; + crossRunning = running; + } + running += weight[i]; + } + if (cross != INT_MAX) + { + atomicMin(&crossSeg, cross); + } + __syncthreads(); + int const seg = crossSeg; + if (seg != INT_MAX && cross == seg) + { + crossBefore = crossRunning; + } + __syncthreads(); + int lastFrom = numSegs - 1; + if (seg != INT_MAX) + { + // One logit of the segment per thread, in index order. + float const* rowLogits = params.logits + static_cast(row) * params.vocabSize; + int const idx = seg * kDirectSegElts + tid; + float const e = idx < params.vocabSize ? directWeight(rowLogits[idx], r, draw.rowMax) : 0.F; + float prefix; + SampleBlockScan(tempStorage.scan).ExclusiveSum(e, prefix); + if (e > 0.F && crossBefore + prefix + e > target) + { + atomicMin(&sampled, idx); + } + __syncthreads(); + if (sampled != INT_MAX) + { + draw.token = sampled; + if (needTokenExp) + { + if (idx == sampled) + { + sampledExp = e; + } + __syncthreads(); + draw.tokenExp = sampledExp; + } + return draw; + } + // The segment's sum passed u * rowSum but its scan did not (rounding): its last token with probability > 0. + lastFrom = seg; + } + draw.token = lastPositiveToken(params, row, r, draw.rowMax, lastFrom, draw.tokenExp, tempStorage); + return draw; +} + +// The per-segment sums of raw weights above pivot0 and above pivot1 for the segments [segBegin, segEnd) of +// `row`, stored in segAbove0 / segAbove1 by the warps of one CTA. A segment whose largest probability is at most +// min(pivot0, pivot1) has none above either pivot (the probability is monotone in the logit) and stores 0 unread. +__device__ void storeAboveSums(DirectKernelParams const& params, int row, DirectRow const& r, float rowMax, + float pivot0, float pivot1, int segBegin, int segEnd) +{ + int const lane = static_cast(threadIdx.x) % kWarpSize; + int64_t const rowSegs = static_cast(row) * params.numSegs; + float const threshold = fminf(pivot0, pivot1); + forEachSegment( + params, row, segBegin, segEnd, [&](float segMaxV) { return __expf(segMaxV - rowMax) > threshold; }, + [&](int seg, float4 const(&x)[kDirectSegVecs], bool const(&valid)[kDirectSegVecs]) + { + float sum0 = 0.F; + float sum1 = 0.F; +#pragma unroll + for (int q = 0; q < kDirectSegVecs; ++q) + { + if (valid[q]) + { + float const xs[kSampleVecSize] = {x[q].x, x[q].y, x[q].z, x[q].w}; +#pragma unroll + for (int j = 0; j < kSampleVecSize; ++j) + { + float const p = directWeight(xs[j], r, rowMax); + sum0 += p > pivot0 ? p : 0.F; + sum1 += p > pivot1 ? p : 0.F; + } + } + } + sum0 = sampleWarpSum(sum0); + sum1 = sampleWarpSum(sum1); + if (lane == 0) + { + params.segAbove0[rowSegs + seg] = sum0; + params.segAbove1[rowSegs + seg] = sum1; + } + }, + [&](int seg) + { + params.segAbove0[rowSegs + seg] = 0.F; + params.segAbove1[rowSegs + seg] = 0.F; + }); +} + +// The sums of segAbove0 and segAbove1 over the row on one CTA, thread t adding its segments [t * perThread, +// (t + 1) * perThread) in order. The sums may come from other CTAs of the row (ld.cg). +template +__device__ void aboveTotals( + DirectKernelParams const& params, int row, float& total0, float& total1, DirectTempStorage& tempStorage) +{ + __shared__ float totals[2]; + int const tid = static_cast(threadIdx.x); + int const numSegs = params.numSegs; + int const perThread = (numSegs + kSampleBlockSize - 1) / kSampleBlockSize; + int64_t const rowSegs = static_cast(row) * numSegs; + float sum0 = 0.F; + float sum1 = 0.F; +#pragma unroll + for (int i = 0; i < kSegsPerThread; ++i) + { + int const seg = tid * perThread + i; + if (i < perThread && seg < numSegs) + { + sum0 += __ldcg(¶ms.segAbove0[rowSegs + seg]); + sum1 += __ldcg(¶ms.segAbove1[rowSegs + seg]); + } + } + float const blockSum0 = SampleBlockReduce(tempStorage.reduce).Sum(sum0); + __syncthreads(); + float const blockSum1 = SampleBlockReduce(tempStorage.reduce).Sum(sum1); + if (tid == 0) + { + totals[0] = blockSum0; + totals[1] = blockSum1; + } + __syncthreads(); + total0 = totals[0]; + total1 = totals[1]; + __syncthreads(); +} + +// One retry inverse-CDF draw on one CTA, for the tokens with raw weight > low (the +// weights of the original draw), from their per-segment sums segSums (stored by the top-p test that +// chose low): the segment where the running sum passes u, then the token among its 512 logits. u at or above the +// sum gives the last token with probability > low. Returns -1 when none has. +template +__device__ int sampleAboveLow(DirectKernelParams const& params, int row, DirectRow const& r, float rowMax, float low, + float const* segSums, float u, DirectTempStorage& tempStorage) +{ + __shared__ int crossSeg; + __shared__ float crossBefore; + __shared__ int lastSeg; + __shared__ int sampledIdx; + int const tid = static_cast(threadIdx.x); + int const numSegs = params.numSegs; + int const perThread = (numSegs + kSampleBlockSize - 1) / kSampleBlockSize; + if (tid == 0) + { + crossSeg = INT_MAX; + sampledIdx = INT_MAX; + } + // Thread t holds segments [t * perThread, (t + 1) * perThread). + float sums[kSegsPerThread]; + float threadSum = 0.F; + int threadLast = -1; +#pragma unroll + for (int i = 0; i < kSegsPerThread; ++i) + { + int const seg = tid * perThread + i; + sums[i] = i < perThread && seg < numSegs ? __ldcg(&segSums[seg]) : 0.F; + threadSum += sums[i]; + if (sums[i] > 0.F) + { + threadLast = seg; + } + } + float before; + SampleBlockScan(tempStorage.scan).ExclusiveSum(threadSum, before); + float running = before; + int cross = INT_MAX; + float crossRunning = 0.F; +#pragma unroll + for (int i = 0; i < kSegsPerThread; ++i) + { + if (cross == INT_MAX && sums[i] > 0.F && running + sums[i] > u) + { + cross = tid * perThread + i; + crossRunning = running; + } + running += sums[i]; + } + __syncthreads(); + int const blockLast = SampleBlockReduceInt(tempStorage.reduceInt).Reduce(threadLast, IntMax{}); + if (cross != INT_MAX) + { + atomicMin(&crossSeg, cross); + } + if (tid == 0) + { + lastSeg = blockLast; + } + __syncthreads(); + int const seg = crossSeg; + if (seg != INT_MAX && cross == seg) + { + crossBefore = crossRunning; + } + // Without a crossing segment (u at or above the sum) the result is the last token with probability > low. + int const scanSeg = seg != INT_MAX ? seg : lastSeg; + __syncthreads(); + if (scanSeg < 0) + { + return -1; + } + // One logit of the segment per thread, in index order. + float const* rowLogits = params.logits + static_cast(row) * params.vocabSize; + int const idx = scanSeg * kDirectSegElts + tid; + float const p = idx < params.vocabSize ? directWeight(rowLogits[idx], r, rowMax) : 0.F; + bool const kept = p > low; + if (seg != INT_MAX) + { + float prefix; + SampleBlockScan(tempStorage.scan).ExclusiveSum(kept ? p : 0.F, prefix); + if (kept && crossBefore + prefix + p > u) + { + atomicMin(&sampledIdx, idx); + } + __syncthreads(); + if (sampledIdx != INT_MAX) + { + int const result = sampledIdx; + __syncthreads(); + return result; + } + // The segment's sum passed u but its scan did not (rounding): its last token with probability > low. + } + int const segLast = SampleBlockReduceInt(tempStorage.reduceInt).Reduce(kept ? idx : -1, IntMax{}); + if (tid == 0) + { + sampledIdx = segLast; + } + __syncthreads(); + int const result = sampledIdx; + __syncthreads(); + return result; +} + +// FULL top-p uses raw weights and a fixed target based on the original filtered mass. +// Strictly greater weight defines the mass above a candidate, preserving complete tie groups. +// Every retry uses the same RNG mapping and support domain as the initial draw. +template +__device__ int directTopPRounds(DirectKernelParams const& params, int row, DirectRow const& r, DirectDraw const& draw, + float pivot0, float pivot1, float aggregate0, float aggregate1, DirectTempStorage& tempStorage) +{ + int token = draw.token; + float const target = r.topP * draw.rowSum; + if (aggregate0 <= target) + { + return token; + } + int64_t const rowSegs = static_cast(row) * params.numSegs; + float low = 0.F; + float high = 1.F; + float q = draw.rowSum; + // low is the first pivot of the last test (its sums are in segAbove0), else its second pivot (segAbove1). + bool lowIsPivot0 = true; + if (aggregate1 <= target) + { + low = pivot0; + high = pivot1; + q = aggregate0; + } + else + { + low = pivot1; + q = aggregate1; + lowIsPivot0 = false; + } + curandStatePhilox4_32_10_t state; + initDirectRowRng(state, params, row); + curand_uniform(&state); + float const* rowLogits = params.logits + static_cast(row) * params.vocabSize; + for (int round = 0; round < 32 && low < high; ++round) + { + float const u = curand_uniform(&state) * q; + float const* lowSums = (lowIsPivot0 ? params.segAbove0 : params.segAbove1) + rowSegs; + token = sampleAboveLow(params, row, r, draw.rowMax, low, lowSums, u, tempStorage); + if (token < 0) + { + return rowArgmax(params, row, r); + } + float const roundPivot0 = directWeight(rowLogits[token], r, draw.rowMax); + float const roundPivot1 = ieeeMul(ieeeAdd(roundPivot0, high), 0.5F); + storeAboveSums(params, row, r, draw.rowMax, roundPivot0, roundPivot1, 0, params.numSegs); + __syncthreads(); + float roundAbove0; + float roundAbove1; + aboveTotals(params, row, roundAbove0, roundAbove1, tempStorage); + if (roundAbove0 <= target) + { + return token; + } + if (roundAbove1 <= target) + { + low = roundPivot0; + high = roundPivot1; + q = roundAbove0; + lowIsPivot0 = true; + } + else + { + low = roundPivot1; + q = roundAbove1; + lowIsPivot0 = false; + } + } + return rowArgmax(params, row, r); +} + +// First direct-sampling kernel, CTA blockIdx.x of row blockIdx.y: max and sum of exp(v - max) of v = logit / T for +// its 16 segments, a warp each. A row with the top-k mask gets its masked partials in the second kernel. +__global__ void __launch_bounds__(kSampleBlockSize) directPartialsKernel(DirectKernelParams params) +{ + int const cta = static_cast(blockIdx.x); + int const row = static_cast(blockIdx.y); + // The temperature is a step input, not written by the preceding kernel. + float const invTemp = inverseTemperature(params.temperatures[row]); +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + // The logits are written by the preceding kernel. Dependents read the outputs only after their own wait. + cudaGridDependencySynchronize(); + cudaTriggerProgrammaticLaunchCompletion(); +#endif + int const seg = cta * kDirectWarps + static_cast(threadIdx.x) / kWarpSize; + if (seg >= params.numSegs) + { + return; + } + float4 x[kDirectSegVecs]; + bool valid[kDirectSegVecs]; + loadSegment(params.logits + static_cast(row) * params.vocabSize, params.vocabSize, seg, + static_cast(threadIdx.x) % kWarpSize, x, valid); + storeSegmentPartials(params, row, seg, invTemp, false, 0.F, x, valid); +} + +// Second direct-sampling kernel, CTA (row blockIdx.x, chunk blockIdx.y): the draw from the segment partials. A row +// with the top-k mask first finds the k-th largest logit and its masked partials. A row with top_p < 1 then runs the +// top-p test and its later rounds; without the mask, its first test is spread over one CTA per chunk of 16 segments +// (each repeats the draw), and the last CTA of the row to finish runs the rounds. Every other row runs on its first +// CTA alone. +template +__global__ void __launch_bounds__(kSampleBlockSize) directDrawKernel(DirectKernelParams params) +{ + __shared__ DirectTempStorage tempStorage; + __shared__ float drawU; + int const row = static_cast(blockIdx.x); + int const chunk = static_cast(blockIdx.y); + DirectRow r = loadDirectRow(params, row); + bool const split = r.nucleus && !r.masked && r.minP == 0.F; + if (chunk != 0 && !split) + { + return; + } + if (threadIdx.x == 0) + { + // The first draw of the row's Philox stream: the seed and the offset are step inputs. + curandStatePhilox4_32_10_t state; + initDirectRowRng(state, params, row); + drawU = curand_uniform(&state); + } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + // The partials are written by the preceding kernel. + cudaGridDependencySynchronize(); + cudaTriggerProgrammaticLaunchCompletion(); +#endif + float pivot = 0.F; + if (r.masked || r.minP > 0.F) + { + float const maximum = originalRowMax(params, row, tempStorage); + if (r.masked) + { + pivot = directPivot(params, row, r.topK); + r.weightFloor = fmaxf(r.minP, __expf(__fmul_rn(pivot, r.invTemp) - maximum)); + } + storeFilteredPartials(params, row, r, maximum); + __syncthreads(); + } + DirectDraw const draw = drawFromSegments(params, row, r, &drawU, r.nucleus, tempStorage); + if (draw.token < 0 || !r.nucleus) + { + int const result = draw.token < 0 ? rowArgmax(params, row, r) : draw.token; + if (chunk == 0 && threadIdx.x == 0) + { + params.tokens[row] = result; + } + return; + } + float const pivot0 = draw.tokenExp; + float const pivot1 = ieeeMul(ieeeAdd(pivot0, 1.F), 0.5F); + if (split) + { + int const segBegin = chunk * kDirectWarps; + storeAboveSums( + params, row, r, draw.rowMax, pivot0, pivot1, segBegin, min(segBegin + kDirectWarps, params.numSegs)); + if (!arriveLastOfRow(¶ms.rowCounters[row], params.numChunks)) + { + return; + } + } + else + { + storeAboveSums(params, row, r, draw.rowMax, pivot0, pivot1, 0, params.numSegs); + __syncthreads(); + } + float aggregate0; + float aggregate1; + aboveTotals(params, row, aggregate0, aggregate1, tempStorage); + int const token + = directTopPRounds(params, row, r, draw, pivot0, pivot1, aggregate0, aggregate1, tempStorage); + if (threadIdx.x == 0) + { + params.tokens[row] = token; + } +} + +} // namespace + +size_t getGlmVerifyHeadSampleWorkspaceSize(int numRows, int vocabSize) +{ + size_t const numSegs = static_cast((vocabSize + kDirectSegElts - 1) / kDirectSegElts); + return ((4 + kWarpSize) * numSegs + kDirectPivotCandidates) * static_cast(numRows) * sizeof(float); +} + +void invokeGlmVerifyHeadSample(GlmVerifyHeadSampleParams const& params) +{ + TLLM_CHECK_WITH_INFO(params.vocabSize == 154880 && params.numRows >= 1 && params.numRows <= 24, + "GLM verify sampling requires 1..24 rows and vocabulary 154880"); + TLLM_CHECK(params.rowCounters != nullptr && params.minP != nullptr && params.topP != nullptr); + int const numSegs = (params.vocabSize + kDirectSegElts - 1) / kDirectSegElts; + int const numChunks = numSampleChunks(params.vocabSize); + size_t const rowSegs = static_cast(params.numRows) * static_cast(numSegs); + auto* base = static_cast(params.workspace); + + DirectKernelParams kernelParams; + kernelParams.logits = params.logits; + kernelParams.temperatures = params.temperatures; + kernelParams.topK = params.topK; + kernelParams.topP = params.topP; + kernelParams.minP = params.minP; + kernelParams.perRowRng = params.perRowRng; + kernelParams.seed = params.seed; + kernelParams.offset = params.offset; + kernelParams.rowCounters = params.rowCounters; + kernelParams.tokens = params.tokens; + kernelParams.segMax = base; + kernelParams.segSum = kernelParams.segMax + rowSegs; + kernelParams.segAbove0 = kernelParams.segSum + rowSegs; + kernelParams.segAbove1 = kernelParams.segAbove0 + rowSegs; + kernelParams.laneMax = kernelParams.segAbove1 + rowSegs; + kernelParams.pivotPool = kernelParams.laneMax + kWarpSize * rowSegs; + kernelParams.vocabSize = params.vocabSize; + kernelParams.numSegs = numSegs; + kernelParams.numChunks = numChunks; + + tensorrt_llm::common::launchWithPdlWhenEnabled("glmVerifyHeadPartials", directPartialsKernel, + dim3(static_cast(numChunks), static_cast(params.numRows)), dim3(kSampleBlockSize), 0, + params.stream, kernelParams); + // With top_p, a row without the top-k mask may spread its first top-p test over numChunks CTAs; the first CTA of + // every row comes first in launch order. + dim3 const drawGrid( + static_cast(params.numRows), static_cast(params.topP != nullptr ? numChunks : 1)); + tensorrt_llm::common::launchWithPdlWhenEnabled( + "glmVerifyHeadDraw", directDrawKernel<1>, drawGrid, dim3(kSampleBlockSize), 0, params.stream, kernelParams); +} + +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmVerifyHeadSample.h b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmVerifyHeadSample.h new file mode 100644 index 000000000000..27ff86509f54 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmVerifyHeadSample.h @@ -0,0 +1,52 @@ +/* + * 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. + */ + +#pragma once +#include "tensorrt_llm/common/config.h" +#include +#include +#include +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa +{ +// Tokens-only FULL sampling for gathered GLM logits. Parameters/RNG arrays are step +// inputs already complete before the preceding model-forward producer. Scratch and +// counters are owned per CUDA stream / graph state and must not be used concurrently. +struct GlmVerifyHeadSampleParams +{ + float const* logits; + float const* temperatures; + int32_t const* topK; + float const* topP; + float const* minP; + int64_t const* seed; + int64_t const* offset; + bool perRowRng; + int32_t* rowCounters; + int32_t* tokens; + void* workspace; + int numRows; + int vocabSize; + cudaStream_t stream; +}; + +size_t getGlmVerifyHeadSampleWorkspaceSize(int numRows, int vocabSize); +void invokeGlmVerifyHeadSample(GlmVerifyHeadSampleParams const& params); +} // namespace kernels::glm_moe_dsa + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index ee3d36f7a8b7..b337923af7fc 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -177,7 +177,7 @@ add_th_common_family( # GLM bindings keep their existing family targets; source lists live with the # model. -add_subdirectory(glmMoeDsa) +add_subdirectory(glmMoeDsaRubin) add_library( th_common SHARED diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/CMakeLists.txt b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/CMakeLists.txt similarity index 76% rename from cpp/tensorrt_llm/thop/glmMoeDsa/CMakeLists.txt rename to cpp/tensorrt_llm/thop/glmMoeDsaRubin/CMakeLists.txt index cad8b5ad52a0..c28c41ca75a7 100644 --- a/cpp/tensorrt_llm/thop/glmMoeDsa/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/CMakeLists.txt @@ -14,7 +14,9 @@ target_sources(th_common_moe PRIVATE glmRouterGemmOp.cpp glmLowMBf16MoeOp.cpp glmRoutedExpertsOp.cpp) target_sources( - th_common_distributed PRIVATE glmDsaHistTopKOp.cpp glmMnnvlAllgatherOp.cpp - glmMnnvlHopOp.cpp glmDraftSampleGlueOp.cpp) + th_common_distributed + PRIVATE glmDsaHistTopKOp.cpp glmMnnvlAllgatherOp.cpp glmMnnvlHopOp.cpp + glmDraftSampleGlueOp.cpp glmIntraNodeReduceOp.cpp) -target_sources(th_common_decoding PRIVATE glmVerifyHeadAcceptOp.cpp) +target_sources(th_common_decoding PRIVATE glmVerifyHeadAcceptOp.cpp + glmVerifyHeadSampleOp.cpp) diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/glmDenseMlpSmallMOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmDenseMlpSmallMOp.cpp similarity index 98% rename from cpp/tensorrt_llm/thop/glmMoeDsa/glmDenseMlpSmallMOp.cpp rename to cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmDenseMlpSmallMOp.cpp index 91b36a423ec4..c0e6ed7b9bb9 100644 --- a/cpp/tensorrt_llm/thop/glmMoeDsa/glmDenseMlpSmallMOp.cpp +++ b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmDenseMlpSmallMOp.cpp @@ -17,7 +17,7 @@ #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/envUtils.h" -#include "tensorrt_llm/kernels/glmMoeDsa/glmDenseMlpSmallM.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmDenseMlpSmallM.h" #include "tensorrt_llm/thop/thUtils.h" #include diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/glmDraftSampleGlueOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmDraftSampleGlueOp.cpp similarity index 99% rename from cpp/tensorrt_llm/thop/glmMoeDsa/glmDraftSampleGlueOp.cpp rename to cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmDraftSampleGlueOp.cpp index b9bbcadff112..4903e1bc1325 100644 --- a/cpp/tensorrt_llm/thop/glmMoeDsa/glmDraftSampleGlueOp.cpp +++ b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmDraftSampleGlueOp.cpp @@ -16,7 +16,7 @@ */ #include "tensorrt_llm/common/mcastDevMemUtils.h" -#include "tensorrt_llm/kernels/glmMoeDsa/glmDraftSampleGlue.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmDraftSampleGlue.h" #include "tensorrt_llm/runtime/mcastDeviceMemory.h" #include "tensorrt_llm/thop/thUtils.h" diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/glmDsaHistTopKOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmDsaHistTopKOp.cpp similarity index 99% rename from cpp/tensorrt_llm/thop/glmMoeDsa/glmDsaHistTopKOp.cpp rename to cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmDsaHistTopKOp.cpp index fb0ab897d100..bc3e549310a7 100644 --- a/cpp/tensorrt_llm/thop/glmMoeDsa/glmDsaHistTopKOp.cpp +++ b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmDsaHistTopKOp.cpp @@ -16,7 +16,7 @@ */ #include "tensorrt_llm/common/mcastDevMemUtils.h" -#include "tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmDsaHistTopK.h" #include "tensorrt_llm/runtime/mcastDeviceMemory.h" #include "tensorrt_llm/thop/thUtils.h" diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/glmFusedAGemmOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmFusedAGemmOp.cpp similarity index 98% rename from cpp/tensorrt_llm/thop/glmMoeDsa/glmFusedAGemmOp.cpp rename to cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmFusedAGemmOp.cpp index 2a33b5fc1e1f..59e402c07495 100644 --- a/cpp/tensorrt_llm/thop/glmMoeDsa/glmFusedAGemmOp.cpp +++ b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmFusedAGemmOp.cpp @@ -16,7 +16,7 @@ */ #include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/kernels/glmMoeDsa/glmFusedAGemm.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmFusedAGemm.h" #include #include diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/glmFusedEmbeddingOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmFusedEmbeddingOp.cpp similarity index 97% rename from cpp/tensorrt_llm/thop/glmMoeDsa/glmFusedEmbeddingOp.cpp rename to cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmFusedEmbeddingOp.cpp index cd95ab8210e7..333b12924d96 100644 --- a/cpp/tensorrt_llm/thop/glmMoeDsa/glmFusedEmbeddingOp.cpp +++ b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmFusedEmbeddingOp.cpp @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/glmMoeDsa/glmFusedEmbeddingLookup.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmFusedEmbeddingLookup.h" #include "tensorrt_llm/thop/thUtils.h" #include diff --git a/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmIntraNodeReduceOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmIntraNodeReduceOp.cpp new file mode 100644 index 000000000000..439d43bf7a78 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmIntraNodeReduceOp.cpp @@ -0,0 +1,217 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#include "tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmIntraNodeReduce.h" +#include "tensorrt_llm/thop/thUtils.h" +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ +namespace +{ +namespace glm = kernels::glm_moe_dsa; +using Pair = std::tuple; +using Quad = std::tuple; + +void bf16(at::Tensor const& t, at::Device device, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.device() == device && t.scalar_type() == at::kBFloat16 && t.is_contiguous() + && reinterpret_cast(t.const_data_ptr()) % 16 == 0, + name, " must be contiguous, 16-byte aligned BF16 on the input CUDA device"); +} + +void validate(at::Tensor const& residual, at::Tensor const& gamma, double epsilon, at::Tensor const& table, + at::Tensor const& flags, at::Tensor const& layout, int64_t rank, int64_t nranks, int64_t bufferBytes, bool finalize) +{ + bf16(residual, residual.device(), "residual"); + TORCH_CHECK(residual.dim() == 2 && residual.size(1) == 6144 && residual.size(0) >= 1 && residual.size(0) <= 24, + "GLM IPC requires BF16 [1..24, 6144] rows"); + bf16(gamma, residual.device(), "gamma"); + TORCH_CHECK(gamma.dim() == 1 && gamma.numel() == 6144, "gamma must have 6144 elements"); + TORCH_CHECK(std::isfinite(epsilon) && epsilon >= 0 && epsilon <= std::numeric_limits::max(), + "invalid RMSNorm epsilon"); + TORCH_CHECK(nranks == 4 && rank >= 0 && rank < nranks, "GLM IPC requires a four-rank in-node TP group"); + TORCH_CHECK(table.device() == residual.device() && table.scalar_type() == at::kLong && table.dim() == 1 + && table.numel() == 3 * nranks + 2 && table.is_contiguous(), + "expected the bound IPC pointer table"); + TORCH_CHECK(flags.device() == residual.device() && flags.scalar_type() == at::kInt && flags.dim() == 1 + && flags.numel() == 3 && flags.is_contiguous(), + "expected bound IPC flags[3]"); + TORCH_CHECK(layout.device() == residual.device() && layout.scalar_type() == at::kLong && layout.dim() == 1 + && layout.numel() == 2 && layout.is_contiguous(), + "expected bound IPC layout[2]"); + TORCH_CHECK( + bufferBytes > 0 && bufferBytes % 16 == 0 && residual.numel() * residual.element_size() * nranks <= bufferBytes, + "rows exceed the bound IPC Lamport buffer capacity"); + TORCH_CHECK(glm::intraNodeReduceSupported(residual.size(0), residual.size(1), nranks, finalize), + "unsupported GLM Rubin IPC launch"); +} + +glm::IntraNodeReduceParams paramsFor(at::Tensor const& residual, at::Tensor const& gamma, double epsilon, + at::Tensor const& table, at::Tensor& flags, at::Tensor& layout, int64_t rank, at::Tensor& normOut, + at::Tensor& residualOut) +{ + glm::IntraNodeReduceParams p{}; + p.workspace = reinterpret_cast(table.const_data_ptr()); + // These are the actual mutated Tensor allocations, not aliases of the pointer table. + p.flags = flags.mutable_data_ptr(); + p.layout = layout.mutable_data_ptr(); + p.residual_in = residual.const_data_ptr(); + p.rms_gamma = gamma.const_data_ptr(); + p.residual_out = residualOut.mutable_data_ptr(); + p.norm_out = normOut.mutable_data_ptr(); + p.rank = static_cast(rank); + p.nranks = 4; + p.size = static_cast(residual.numel()); + p.hidden_dim = 6144; + p.rms_eps = static_cast(epsilon); + p.stream = at::cuda::getCurrentCUDAStream(residual.get_device()); + return p; +} +} // namespace + +Pair glmIntraNodeReduce(at::Tensor const& input, at::Tensor const& residual, at::Tensor const& gamma, double epsilon, + at::Tensor const& table, at::Tensor& flags, at::Tensor& layout, int64_t rank, int64_t nranks, int64_t bufferBytes, + bool triggerCompletionAtEnd) +{ + c10::cuda::CUDAGuard guard(residual.device()); + validate(residual, gamma, epsilon, table, flags, layout, rank, nranks, bufferBytes, false); + bf16(input, residual.device(), "input"); + TORCH_CHECK(input.sizes() == residual.sizes(), "input and residual shapes must match"); + auto normOut = at::empty_like(residual); + auto residualOut = at::empty_like(residual); + auto p = paramsFor(residual, gamma, epsilon, table, flags, layout, rank, normOut, residualOut); + p.allreduce_in = input.const_data_ptr(); + p.trigger_completion_at_end = triggerCompletionAtEnd; + glm::intraNodeReduce(p); + return {normOut, residualOut}; +} + +Pair glmIntraNodeMoeFinalize(at::Tensor const& fc2, at::Tensor const& weights, at::Tensor const& indices, + at::Tensor const& shared, at::Tensor const& residual, at::Tensor const& gamma, double epsilon, + at::Tensor const& table, at::Tensor& flags, at::Tensor& layout, int64_t rank, int64_t nranks, int64_t bufferBytes) +{ + c10::cuda::CUDAGuard guard(residual.device()); + validate(residual, gamma, epsilon, table, flags, layout, rank, nranks, bufferBytes, true); + bf16(fc2, residual.device(), "fc2"); + bf16(shared, residual.device(), "shared"); + TORCH_CHECK(fc2.dim() == 2 && fc2.size(1) == 6144 && shared.sizes() == residual.sizes(), + "FC2 must have unpadded hidden6144; shared output must match residual"); + TORCH_CHECK(indices.device() == residual.device() && indices.scalar_type() == at::kInt && indices.is_contiguous() + && indices.dim() == 2 && indices.size(0) == residual.size(0) && indices.size(1) >= 1 + && indices.size(1) <= 8, + "indices must be int32 [rows, 1..8]"); + TORCH_CHECK(weights.device() == residual.device() && weights.scalar_type() == at::kBFloat16 + && weights.is_contiguous() && weights.sizes() == indices.sizes(), + "weights must be BF16 matching indices"); + auto normOut = at::empty_like(residual); + auto residualOut = at::empty_like(residual); + auto p = paramsFor(residual, gamma, epsilon, table, flags, layout, rank, normOut, residualOut); + p.allreduce_in = fc2.const_data_ptr(); + p.expert_scale_factor = weights.const_data_ptr(); + p.expanded_idx_to_permuted_idx = indices.const_data_ptr(); + p.shared_expert_output = shared.const_data_ptr(); + p.top_k = static_cast(indices.size(1)); + glm::intraNodeMoeFinalize(p); + return {normOut, residualOut}; +} + +// The table must be bound once to these real flags/layout allocations by _ipc_workspace. +// The shared FP4 kernel accesses the same state indirectly via the table's last two pointers. +Quad glmPreMoeFp4Linear(at::Tensor const& input, at::Tensor const& residual, at::Tensor const& gamma, double epsilon, + at::Tensor const& scale, at::Tensor const& table, at::Tensor& flags, at::Tensor& layout, int64_t rank, + int64_t nranks, int64_t bufferBytes) +{ + c10::cuda::CUDAGuard guard(residual.device()); + validate(residual, gamma, epsilon, table, flags, layout, rank, nranks, bufferBytes, false); + bf16(input, residual.device(), "input"); + TORCH_CHECK(input.sizes() == residual.sizes(), "input and residual shapes must match"); + TORCH_CHECK(scale.device() == residual.device() && scale.scalar_type() == at::kFloat && scale.numel() == 1 + && scale.is_contiguous(), + "scale must be one contiguous CUDA FP32 value"); + auto normOut = at::empty_like(input); + auto residualOut = at::empty_like(residual); + auto packed = at::empty({input.size(0), 3072}, input.options().dtype(at::kByte)); + auto scales = at::empty({input.size(0), 384}, input.options().dtype(at::kByte)); + kernels::ar_fusion::AllReduceFusionParams p{}; + p.nranks = 4; + p.rank = static_cast(rank); + p.dtype = DataType::kBF16; + p.size = static_cast(input.numel()); + p.hidden_dim = 6144; + p.workspace = reinterpret_cast(const_cast(table.const_data_ptr())); + p.allreduce_in = const_cast(input.const_data_ptr()); + p.residual_in = const_cast(residual.const_data_ptr()); + p.rms_gamma = const_cast(gamma.const_data_ptr()); + p.scale_factor = const_cast(scale.const_data_ptr()); + p.norm_out = normOut.mutable_data_ptr(); + p.residual_out = residualOut.mutable_data_ptr(); + p.quant_out = packed.mutable_data_ptr(); + p.scale_out = scales.mutable_data_ptr(); + p.rms_eps = static_cast(epsilon); + p.use_oneshot = true; + p.layout = QuantizationSFLayout::LINEAR; + p.pattern = kernels::ar_fusion::AllReduceFusionPattern::kARResidualRMSNormOutFP4Quant; + p.trigger_completion_at_end = false; + p.stream = at::cuda::getCurrentCUDAStream(input.get_device()); + kernels::ar_fusion::allreduce_fusion_op(p); + return {normOut, residualOut, packed, scales}; +} + +bool glmIntraNodeReduceSupported(int64_t rows, int64_t hidden, int64_t nranks, bool finalize) +{ + return rows >= 1 && rows <= 24 && hidden == 6144 && nranks == 4 + && glm::intraNodeReduceSupported(static_cast(rows), 6144, 4, finalize); +} + +void rejectGlmIpcFunctionalization(c10::OperatorHandle const&, torch::jit::Stack*) +{ + TORCH_CHECK(false, + "GLM IPC collectives support eager and ordinary CUDA graphs only: " + "functionalization cannot clone the peer allocations referenced by the IPC pointer table"); +} +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_intra_node_reduce(Tensor input, Tensor residual, Tensor gamma, float epsilon, Tensor table, " + "Tensor(a!) flags, Tensor(b!) layout, int rank, int nranks, int buffer_bytes, bool trigger_completion_at_end) " + "-> (Tensor, Tensor)"); + m.def( + "glm_intra_node_moe_finalize(Tensor fc2, Tensor weights, Tensor indices, Tensor shared, Tensor residual, " + "Tensor gamma, float epsilon, Tensor table, Tensor(a!) flags, Tensor(b!) layout, int rank, int nranks, " + "int buffer_bytes) -> (Tensor, Tensor)"); + m.def( + "glm_pre_moe_fp4_linear(Tensor input, Tensor residual, Tensor gamma, float epsilon, Tensor scale, Tensor " + "table, " + "Tensor(a!) flags, Tensor(b!) layout, int rank, int nranks, int buffer_bytes) -> (Tensor, Tensor, Tensor, " + "Tensor)"); + m.def("glm_intra_node_reduce_supported", &tensorrt_llm::torch_ext::glmIntraNodeReduceSupported); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_intra_node_reduce", &tensorrt_llm::torch_ext::glmIntraNodeReduce); + m.impl("glm_intra_node_moe_finalize", &tensorrt_llm::torch_ext::glmIntraNodeMoeFinalize); + m.impl("glm_pre_moe_fp4_linear", &tensorrt_llm::torch_ext::glmPreMoeFp4Linear); +} + +TORCH_LIBRARY_IMPL(trtllm, Functionalize, m) +{ + m.impl("glm_intra_node_reduce", + torch::CppFunction::makeFromBoxedFunction<&tensorrt_llm::torch_ext::rejectGlmIpcFunctionalization>()); + m.impl("glm_intra_node_moe_finalize", + torch::CppFunction::makeFromBoxedFunction<&tensorrt_llm::torch_ext::rejectGlmIpcFunctionalization>()); + m.impl("glm_pre_moe_fp4_linear", + torch::CppFunction::makeFromBoxedFunction<&tensorrt_llm::torch_ext::rejectGlmIpcFunctionalization>()); +} diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/glmLowMBf16MoeOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmLowMBf16MoeOp.cpp similarity index 61% rename from cpp/tensorrt_llm/thop/glmMoeDsa/glmLowMBf16MoeOp.cpp rename to cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmLowMBf16MoeOp.cpp index e414e334d7b8..092209a92fb7 100644 --- a/cpp/tensorrt_llm/thop/glmMoeDsa/glmLowMBf16MoeOp.cpp +++ b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmLowMBf16MoeOp.cpp @@ -16,7 +16,7 @@ */ #include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/kernels/glmMoeDsa/glmLowMBf16Moe.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmLowMBf16Moe.h" #include "tensorrt_llm/thop/thUtils.h" #include @@ -36,10 +36,10 @@ bool isAligned16(at::Tensor const& t) } } // namespace -// One token through an MoE layer with unquantized BF16 SwiGLU experts held as TP slices and one-group DeepSeek-V3 -// routing: routing from the router logits, FC1 + SwiGLU, and FC2 + the weighted sum over the routed experts, in two -// kernels (glmLowMBf16Moe.h). The weights are read in the layout CutlassFusedMoE keeps: w3_w1 [E, 2I, H] with the up -// (W3) rows first, w2 [E, H, I]. Returns [1, H]. +// One to six tokens through an MoE layer with unquantized BF16 SwiGLU experts held as TP slices and one-group +// DeepSeek-V3 routing: routing from the router logits, FC1 + SwiGLU, and FC2 + the weighted sum over the routed +// experts, in two kernels (glmLowMBf16Moe.h). The weights are read in the layout CutlassFusedMoE keeps: w3_w1 [E, 2I, +// H] with the up (W3) rows first, w2 [E, H, I]. Returns [M, H]. at::Tensor glmLowmBf16Moe(at::Tensor const& x, at::Tensor const& routerLogits, at::Tensor const& eScoreCorrectionBias, at::Tensor const& w3W1, at::Tensor const& w2, int64_t topK, double routedScalingFactor) { @@ -60,7 +60,9 @@ at::Tensor glmLowmBf16Moe(at::Tensor const& x, at::Tensor const& routerLogits, a } TORCH_CHECK(tensorrt_llm::common::getSMVersion() >= 90, "glm_lowm_bf16_moe needs SM90 or newer"); - TORCH_CHECK(x.dim() == 2 && x.size(0) == 1, "glm_lowm_bf16_moe takes one token, got x of shape ", x.sizes()); + TORCH_CHECK(x.dim() == 2 && x.size(0) >= 1 && x.size(0) <= tk::kLowMBf16MoeMaxTokens, + "glm_lowm_bf16_moe takes 1 to ", tk::kLowMBf16MoeMaxTokens, " tokens, got ", x.sizes()); + int64_t const numTokens = x.size(0); TORCH_CHECK(w3W1.dim() == 3 && w2.dim() == 3, "glm_lowm_bf16_moe expects 3D expert weights"); int64_t const hiddenSize = x.size(1); int64_t const numExperts = w3W1.size(0); @@ -73,17 +75,18 @@ at::Tensor glmLowmBf16Moe(at::Tensor const& x, at::Tensor const& routerLogits, a w3W1.sizes(), " does not match x ", x.sizes(), " and w2 ", w2.sizes()); TORCH_CHECK(w2.size(0) == numExperts && w2.size(1) == hiddenSize, "glm_lowm_bf16_moe: w2 of shape ", w2.sizes(), " does not match x ", x.sizes(), " and w3_w1 ", w3W1.sizes()); - TORCH_CHECK(routerLogits.dim() == 2 && routerLogits.size(0) == 1 && routerLogits.size(1) == numExperts, - "glm_lowm_bf16_moe: router_logits must be [1, ", numExperts, "], got ", routerLogits.sizes()); + TORCH_CHECK(routerLogits.dim() == 2 && routerLogits.size(0) == numTokens && routerLogits.size(1) == numExperts, + "glm_lowm_bf16_moe: router_logits must be [num_tokens, ", numExperts, "], got ", routerLogits.sizes()); TORCH_CHECK(eScoreCorrectionBias.dim() == 1 && eScoreCorrectionBias.size(0) == numExperts, "glm_lowm_bf16_moe: e_score_correction_bias must be [", numExperts, "], got ", eScoreCorrectionBias.sizes()); - TORCH_CHECK(tk::lowMBf16MoeSupportsShape(hiddenSize, interSize), "glm_lowm_bf16_moe: unsupported hidden size ", - hiddenSize, " or intermediate size ", interSize); + TORCH_CHECK(numTokens == 1 ? tk::lowMBf16MoeSupportsShape(hiddenSize, interSize) + : tk::lowMBf16MoeMultiSupportsShape(numTokens, hiddenSize, interSize), + "glm_lowm_bf16_moe: unsupported hidden size ", hiddenSize, " or intermediate size ", interSize); TORCH_CHECK(isAligned16(x) && isAligned16(w3W1) && isAligned16(w2), "glm_lowm_bf16_moe needs 16-byte aligned x, w3_w1 and w2"); - at::Tensor out = at::empty({1, hiddenSize}, x.options()); - at::Tensor act = at::empty({topK, interSize}, x.options()); + at::Tensor out = at::empty({numTokens, hiddenSize}, x.options()); + at::Tensor act = at::empty({numTokens * topK, interSize}, x.options()); auto const stream = at::cuda::getCurrentCUDAStream(x.get_device()); auto* outPtr = static_cast<__nv_bfloat16*>(out.data_ptr()); auto* actPtr = static_cast<__nv_bfloat16*>(act.data_ptr()); @@ -91,17 +94,36 @@ at::Tensor glmLowmBf16Moe(at::Tensor const& x, at::Tensor const& routerLogits, a auto const* logitsPtr = static_cast(routerLogits.data_ptr()); auto const* w3W1Ptr = static_cast<__nv_bfloat16 const*>(w3W1.data_ptr()); auto const* w2Ptr = static_cast<__nv_bfloat16 const*>(w2.data_ptr()); - if (eScoreCorrectionBias.scalar_type() == at::ScalarType::Float) + if (numTokens == 1) { - tk::invokeLowMBf16Moe(outPtr, actPtr, xPtr, logitsPtr, - static_cast(eScoreCorrectionBias.data_ptr()), w3W1Ptr, w2Ptr, static_cast(hiddenSize), - static_cast(interSize), routedScalingFactor, stream); + if (eScoreCorrectionBias.scalar_type() == at::ScalarType::Float) + { + tk::invokeLowMBf16Moe(outPtr, actPtr, xPtr, logitsPtr, + static_cast(eScoreCorrectionBias.data_ptr()), w3W1Ptr, w2Ptr, + static_cast(hiddenSize), static_cast(interSize), routedScalingFactor, stream); + } + else + { + tk::invokeLowMBf16Moe<__nv_bfloat16>(outPtr, actPtr, xPtr, logitsPtr, + static_cast<__nv_bfloat16 const*>(eScoreCorrectionBias.data_ptr()), w3W1Ptr, w2Ptr, + static_cast(hiddenSize), static_cast(interSize), routedScalingFactor, stream); + } } else { - tk::invokeLowMBf16Moe<__nv_bfloat16>(outPtr, actPtr, xPtr, logitsPtr, - static_cast<__nv_bfloat16 const*>(eScoreCorrectionBias.data_ptr()), w3W1Ptr, w2Ptr, - static_cast(hiddenSize), static_cast(interSize), routedScalingFactor, stream); + if (eScoreCorrectionBias.scalar_type() == at::ScalarType::Float) + { + tk::invokeLowMBf16MoeMulti(outPtr, actPtr, xPtr, logitsPtr, + static_cast(eScoreCorrectionBias.data_ptr()), w3W1Ptr, w2Ptr, static_cast(numTokens), + static_cast(hiddenSize), static_cast(interSize), routedScalingFactor, stream); + } + else + { + tk::invokeLowMBf16MoeMulti<__nv_bfloat16>(outPtr, actPtr, xPtr, logitsPtr, + static_cast<__nv_bfloat16 const*>(eScoreCorrectionBias.data_ptr()), w3W1Ptr, w2Ptr, + static_cast(numTokens), static_cast(hiddenSize), static_cast(interSize), + routedScalingFactor, stream); + } } return out; } diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/glmMlaKvHalfAppendOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmMlaKvHalfAppendOp.cpp similarity index 99% rename from cpp/tensorrt_llm/thop/glmMoeDsa/glmMlaKvHalfAppendOp.cpp rename to cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmMlaKvHalfAppendOp.cpp index c2d2759a08f4..04138f3fc84a 100644 --- a/cpp/tensorrt_llm/thop/glmMoeDsa/glmMlaKvHalfAppendOp.cpp +++ b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmMlaKvHalfAppendOp.cpp @@ -17,7 +17,7 @@ #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/kernels/glmMoeDsa/glmMlaKvHalfAppend.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmMlaKvHalfAppend.h" #include "tensorrt_llm/runtime/torchUtils.h" #include "tensorrt_llm/thop/attentionOp.h" diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/glmMnnvlAllgatherOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmMnnvlAllgatherOp.cpp similarity index 97% rename from cpp/tensorrt_llm/thop/glmMoeDsa/glmMnnvlAllgatherOp.cpp rename to cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmMnnvlAllgatherOp.cpp index ff992084c3d2..966b75b40c82 100644 --- a/cpp/tensorrt_llm/thop/glmMoeDsa/glmMnnvlAllgatherOp.cpp +++ b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmMnnvlAllgatherOp.cpp @@ -16,7 +16,7 @@ */ #include "tensorrt_llm/common/mcastDevMemUtils.h" -#include "tensorrt_llm/kernels/glmMoeDsa/glmMnnvlAllgather.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmMnnvlAllgather.h" #include "tensorrt_llm/runtime/mcastDeviceMemory.h" #include "tensorrt_llm/thop/thUtils.h" @@ -32,7 +32,7 @@ namespace torch_ext // All-gather of BF16 rows along their last dimension, widened to FP32, over a dedicated MNNVL exchange workspace // (comm_buffer, the unicast view of its three Lamport buffers, and buffer_flags): the rank-ordered concatenation of -// the shards (kernels/glmMoeDsa/glmMnnvlAllgather.h). +// the shards (kernels/glmMoeDsaRubin/glmMnnvlAllgather.h). at::Tensor glmMnnvlAllgather( at::Tensor const& input, at::Tensor& comm_buffer, at::Tensor& buffer_flags, int64_t group_size) { diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/glmMnnvlHopOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmMnnvlHopOp.cpp similarity index 96% rename from cpp/tensorrt_llm/thop/glmMoeDsa/glmMnnvlHopOp.cpp rename to cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmMnnvlHopOp.cpp index 2504d971622d..3571cfff0ff3 100644 --- a/cpp/tensorrt_llm/thop/glmMoeDsa/glmMnnvlHopOp.cpp +++ b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmMnnvlHopOp.cpp @@ -16,7 +16,7 @@ */ #include "tensorrt_llm/common/mcastDevMemUtils.h" -#include "tensorrt_llm/kernels/glmMoeDsa/glmMnnvlHop.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmMnnvlHop.h" #include "tensorrt_llm/runtime/mcastDeviceMemory.h" #include "tensorrt_llm/thop/thUtils.h" @@ -111,10 +111,10 @@ glm::MnnvlHopParams makeHopParams(tensorrt_llm::runtime::McastDeviceMemory* mcas } // namespace -// The one-shot MNNVL all-reduce of BF16 rows with the residual add and RMSNorm (kernels/glmMoeDsa/glmMnnvlHop.h) on -// the MNNVL all-reduce workspace comm_buffer / buffer_flags. Returns [norm_out, residual_out]. With shard_published the -// producer of input already stored this rank's rows into the workspace (the o_proj GEMM's Lamport publish), and input -// only gives the shape. +// The one-shot MNNVL all-reduce of BF16 rows with the residual add and RMSNorm (kernels/glmMoeDsaRubin/glmMnnvlHop.h) +// on the MNNVL all-reduce workspace comm_buffer / buffer_flags. Returns [norm_out, residual_out]. With shard_published +// the producer of input already stored this rank's rows into the workspace (the o_proj GEMM's Lamport publish), and +// input only gives the shape. std::vector glmMnnvlHop(at::Tensor const& input, at::Tensor const& residual, at::Tensor const& gamma, double epsilon, at::Tensor& comm_buffer, at::Tensor& buffer_flags, bool shard_published) { diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/glmRoutedExpertsOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmRoutedExpertsOp.cpp similarity index 94% rename from cpp/tensorrt_llm/thop/glmMoeDsa/glmRoutedExpertsOp.cpp rename to cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmRoutedExpertsOp.cpp index cedc6b8f05cd..0fbb447ed548 100644 --- a/cpp/tensorrt_llm/thop/glmMoeDsa/glmRoutedExpertsOp.cpp +++ b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmRoutedExpertsOp.cpp @@ -17,8 +17,9 @@ #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/envUtils.h" -#include "tensorrt_llm/kernels/glmMoeDsa/glmMoeFusedRouter.h" -#include "tensorrt_llm/kernels/glmMoeDsa/glmMoeGemm2K128.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeFusedRouter.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeGemm2K128.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmMoeGemm2K256.h" #include "tensorrt_llm/kernels/moe/trtllmGen/DevKernel.h" #include "tensorrt_llm/kernels/moe/trtllmGen/runner.h" #include "tensorrt_llm/kernels/quantization.h" @@ -49,7 +50,7 @@ namespace glmKernels = tensorrt_llm::kernels::glm_moe_dsa; using trtllmGenMoe::Routing::RoutingMethodType; // FC1 and FC2 of GLM-5's routed experts for one token tile size: the TRTLLM-Gen permute + FC1 + SwiGLU + NVFP4 -// requantization kernel, then FC2 on the TRTLLM-Gen kernel or on glmMoeGemm2K128 where it covers the call. A config +// requantization kernel, then FC2 on TRTLLM-Gen or the owned K128/K256 kernel where it covers the call. A config // is a pair of passing FC1 and FC2 kernel configs, enumerated as the TRTLLM-Gen MoE runner enumerates them. class RoutedExpertsGemms { @@ -115,7 +116,8 @@ class RoutedExpertsGemms // hiddenStatesReady: an event recorded once hidden_states and their scales are written, or null. FC1, their first // reader, waits on it, so the routing launched before FC1 does not. void run(trtllmGenMoe::MoE::MoERunnerArgs const& args, trtllmGenMoe::MoE::MoEWorkspace const& workspace, - cudaEvent_t hiddenStatesReady, bool preferGemm2K128, int device, cudaStream_t stream, int64_t configIndex) + cudaEvent_t hiddenStatesReady, bool preferSpecializedGemm2, int device, cudaStream_t stream, + int64_t configIndex) { sync_check_cuda_error(stream); auto const& config = mConfigs[configIndex]; @@ -140,7 +142,7 @@ class RoutedExpertsGemms args.num_tokens, args.top_k, args.local_num_experts, mTileTokensDim); // glmMoeGemm2K128 covers FC2 at intermediate size 128 and tile N 8 (GLM-5 decode at MoE TP16), with the output // rows of the TRTLLM-Gen FC2 kernel, bitwise. - if (preferGemm2K128 && mTileTokensDim == 8 && args.intermediate_size == 128 + if (preferSpecializedGemm2 && mTileTokensDim == 8 && args.intermediate_size == 128 && validIntermediateSize == args.intermediate_size && validHiddenSize == gemm2HiddenSize && args.gemm2_bias == nullptr && args.output2_scales_scalar != nullptr && glmKernels::gemm2K128::isSupported(gemm2HiddenSize, gemm2MaxNumCtas)) @@ -167,6 +169,29 @@ class RoutedExpertsGemms gemm2Data.sharedMemCarveout = static_cast(cudaSharedmemCarveoutMaxShared); glmKernels::gemm2K128::run(gemm2Data, stream); } + else if (preferSpecializedGemm2 && mTileTokensDim == 8 && args.intermediate_size == 256 + && validIntermediateSize == args.intermediate_size && validHiddenSize == gemm2HiddenSize + && args.gemm2_bias == nullptr && args.output2_scales_scalar != nullptr + && glmKernels::gemm2K256::isSupported(gemm2HiddenSize, gemm2MaxNumCtas)) + { + glmKernels::gemm2K256::Data gemm2Data{}; + gemm2Data.weights = args.gemm2_weights; + gemm2Data.weightsScale = args.gemm2_weights_scale; + gemm2Data.act = workspace.gemm1_output; + gemm2Data.actScale = workspace.gemm1_output_scale; + gemm2Data.outputScale = args.output2_scales_scalar; + gemm2Data.numNonExitingCtas = workspace.num_non_exiting_ctas; + gemm2Data.ctaIdxXyToBatchIdx = workspace.cta_idx_xy_to_batch_idx; + gemm2Data.ctaIdxXyToMnLimit = workspace.cta_idx_xy_to_mn_limit; + gemm2Data.output = workspace.gemm2_output; + gemm2Data.hiddenSize = gemm2HiddenSize; + gemm2Data.numLocalExperts = args.local_num_experts; + gemm2Data.maxNumCtas = gemm2MaxNumCtas; + gemm2Data.usePdl = tensorrt_llm::common::getEnvEnablePDL(); + // Keep the largest shared-memory partition available for the next layer's CTAs. + gemm2Data.sharedMemCarveout = static_cast(cudaSharedmemCarveoutMaxShared); + glmKernels::gemm2K256::run(gemm2Data, stream); + } else { mGemm2.run(workspace.gemm1_output, workspace.gemm1_output_scale, args.gemm2_weights, @@ -225,7 +250,7 @@ std::vector runRoutedExperts(torch::optional const torch::optional const& topk_weights, torch::optional const& topk_ids, torch::optional const& out_tensor, std::optional const hidden_states_ready_event, torch::optional const& router_hidden_states, torch::optional const& router_weight, - torch::optional const& router_workspace, bool const prefer_gemm2_k128) + torch::optional const& router_workspace, bool const prefer_specialized_gemm2) { TORCH_CHECK(tensorrt_llm::common::isSM100Family(), "GLM routed experts run on SM100-family GPUs only."); TORCH_CHECK(tile_tokens_dim == 8 || tile_tokens_dim == 16 || tile_tokens_dim == 32 || tile_tokens_dim == 64 @@ -560,8 +585,8 @@ std::vector runRoutedExperts(torch::optional const ? reinterpret_cast(static_cast(hidden_states_ready_event.value())) : nullptr; auto const& moe_stream = at::cuda::getCurrentCUDAStream(hidden_states.get_device()); - gemms.run( - args, workspace, hidden_states_ready, prefer_gemm2_k128, hidden_states.get_device(), moe_stream, configIndex); + gemms.run(args, workspace, hidden_states_ready, prefer_specialized_gemm2, hidden_states.get_device(), moe_stream, + configIndex); if (!do_finalize) { @@ -624,7 +649,7 @@ class GlmRoutedExpertsRunner : public torch::CustomClassHolder torch::optional const& topk_ids, torch::optional const& output, std::optional const hidden_states_ready_event, torch::optional const& router_hidden_states, torch::optional const& router_weight, - torch::optional const& router_workspace, bool const prefer_gemm2_k128) + torch::optional const& router_workspace, bool const prefer_specialized_gemm2) { int64_t tileN = moeConfigIndex[0]; int64_t config = moeConfigIndex[1]; @@ -644,7 +669,7 @@ class GlmRoutedExpertsRunner : public torch::CustomClassHolder num_experts, top_k, n_group, topk_group, intermediate_size, local_expert_offset, local_num_experts, routed_scaling_factor, tileN, routing_method_type, do_finalize, getGemms(tileN), config, topk_weights, topk_ids, output, hidden_states_ready_event, router_hidden_states, router_weight, router_workspace, - prefer_gemm2_k128); + prefer_specialized_gemm2); } private: diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/glmRouterGemmOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmRouterGemmOp.cpp similarity index 96% rename from cpp/tensorrt_llm/thop/glmMoeDsa/glmRouterGemmOp.cpp rename to cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmRouterGemmOp.cpp index cbaa789904ce..7817a500f31c 100644 --- a/cpp/tensorrt_llm/thop/glmMoeDsa/glmRouterGemmOp.cpp +++ b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmRouterGemmOp.cpp @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/glmMoeDsa/glmRouterGemm.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmRouterGemm.h" #include #include @@ -42,9 +42,9 @@ bool isAligned(torch::Tensor const& tensor) } // namespace -//! The GLM-5 MoE router logits for up to 16 tokens. +//! The GLM-5 MoE router logits for up to 32 tokens. //! -//! @param hiddenStates BF16 CUDA tensor [num_tokens, 6144], contiguous, num_tokens in [1, 16]. +//! @param hiddenStates BF16 CUDA tensor [num_tokens, 6144], contiguous, num_tokens in [1, 32]. //! @param weight BF16 CUDA tensor [256, 6144], contiguous: the router (gate) weight as stored. //! @return FP32 tensor [num_tokens, 256] = hiddenStates @ weight^T, FP32 accumulation. torch::Tensor glmRouterGemm(torch::Tensor const& hiddenStates, torch::Tensor const& weight) diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/glmSparseMlaFlashDecodeOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmSparseMlaFlashDecodeOp.cpp similarity index 51% rename from cpp/tensorrt_llm/thop/glmMoeDsa/glmSparseMlaFlashDecodeOp.cpp rename to cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmSparseMlaFlashDecodeOp.cpp index 019d75820aa9..7aa8bdbb1db4 100644 --- a/cpp/tensorrt_llm/thop/glmMoeDsa/glmSparseMlaFlashDecodeOp.cpp +++ b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmSparseMlaFlashDecodeOp.cpp @@ -16,7 +16,7 @@ */ #include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/kernels/glmMoeDsa/glmSparseMlaFlashDecode.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmSparseMlaFlashDecode.h" #include "tensorrt_llm/runtime/torchUtils.h" #include @@ -53,7 +53,7 @@ bool isBulkCopyAligned(th::Tensor const& t) } // namespace -// Generation sparse MLA of GLM-5.2 at 4 query heads per rank, with the v_b BMM folded in: writes +// Generation sparse MLA of GLM-5.2 at 4 or 8 query heads per rank, with the v_b BMM folded in: writes // output[t][h * 256 + v] = sum_d bf16(attention(q[t, h]))[d] * v_b_proj[h][v][d] for the first q.size(0) rows of // output. With block_table, indices hold token positions that the kernel maps to pool rows as // convert_req_index_to_global does (block_table[s][p / tokens_per_block] * page_index_scale * tokens_per_block @@ -75,10 +75,10 @@ void glmSparseMlaFlashDecode(th::Tensor const& q, th::Tensor const& kvPool, th:: "glm_sparse_mla_flash_decode takes tensors on one device"); TORCH_CHECK(isE4m3Storage(q), "q must hold E4M3 values (float8_e4m3fn or uint8)"); - TORCH_CHECK(q.dim() == 3 && q.size(1) == tk::kGlmSparseMlaNumHeads && q.size(2) == tk::kGlmSparseMlaHeadDimQk - && q.is_contiguous() && isBulkCopyAligned(q), - "q must be a contiguous, 16-byte aligned [tokens, ", tk::kGlmSparseMlaNumHeads, ", ", - tk::kGlmSparseMlaHeadDimQk, "] tensor"); + TORCH_CHECK(q.dim() == 3 && (q.size(1) == tk::kGlmSparseMlaNumHeads || q.size(1) == tk::kGlmSparseMlaNumHeadsTp8) + && q.size(2) == tk::kGlmSparseMlaHeadDimQk && q.is_contiguous() && isBulkCopyAligned(q), + "q must be a contiguous, 16-byte aligned [tokens, 4 or 8, ", tk::kGlmSparseMlaHeadDimQk, "] tensor"); + int64_t const numHeads = q.size(1); int64_t const numTokens = q.size(0); TORCH_CHECK(numTokens <= kMaxQueryTokens, "glm_sparse_mla_flash_decode takes at most ", kMaxQueryTokens, " query tokens, got ", numTokens); @@ -105,15 +105,14 @@ void glmSparseMlaFlashDecode(th::Tensor const& q, th::Tensor const& kvPool, th:: TORCH_CHECK(mlaBmm2Scale.scalar_type() == th::kFloat32 && mlaBmm2Scale.numel() >= 1, "mla_bmm2_scale must hold a float32 value"); - TORCH_CHECK(vBProj.scalar_type() == th::kBFloat16 && vBProj.dim() == 3 - && vBProj.size(0) == tk::kGlmSparseMlaNumHeads && vBProj.size(1) == tk::kGlmSparseMlaVHeadDim - && vBProj.size(2) == tk::kGlmSparseMlaHeadDimV && vBProj.is_contiguous() && isBulkCopyAligned(vBProj), - "v_b_proj must be a contiguous, 16-byte aligned BF16 [", tk::kGlmSparseMlaNumHeads, ", ", - tk::kGlmSparseMlaVHeadDim, ", ", tk::kGlmSparseMlaHeadDimV, "] tensor"); + TORCH_CHECK(vBProj.scalar_type() == th::kBFloat16 && vBProj.dim() == 3 && vBProj.size(0) == numHeads + && vBProj.size(1) == tk::kGlmSparseMlaVHeadDim && vBProj.size(2) == tk::kGlmSparseMlaHeadDimV + && vBProj.is_contiguous() && isBulkCopyAligned(vBProj), + "v_b_proj must be a contiguous, 16-byte aligned BF16 [", numHeads, ", ", tk::kGlmSparseMlaVHeadDim, ", ", + tk::kGlmSparseMlaHeadDimV, "] tensor"); TORCH_CHECK(output.scalar_type() == th::kBFloat16 && output.dim() == 2 && output.size(0) >= numTokens - && output.size(1) == tk::kGlmSparseMlaNumHeads * tk::kGlmSparseMlaVHeadDim && output.stride(1) == 1, - "output must be a BF16 [>= tokens, ", tk::kGlmSparseMlaNumHeads * tk::kGlmSparseMlaVHeadDim, - "] tensor with unit column stride"); + && output.size(1) == numHeads * tk::kGlmSparseMlaVHeadDim && output.stride(1) == 1, + "output must be a BF16 [>= tokens, ", numHeads * tk::kGlmSparseMlaVHeadDim, "] tensor with unit column stride"); tk::GlmSparseMlaFlashDecodeParams params{}; params.q = q.data_ptr(); @@ -125,6 +124,7 @@ void glmSparseMlaFlashDecode(th::Tensor const& q, th::Tensor const& kvPool, th:: params.vbProj = vBProj.data_ptr(); params.out = output.data_ptr(); params.outStride = output.stride(0); + params.numHeadsQ = static_cast(numHeads); params.numTokensQ = static_cast(numTokens); params.seqLenQ = static_cast(seqLenQ); params.indicesStride = static_cast(indices.stride(0)); @@ -159,6 +159,104 @@ void glmSparseMlaFlashDecode(th::Tensor const& q, th::Tensor const& kvPool, th:: tk::invokeGlmSparseMlaFlashDecode(params, stream); } +// Separate latent ABI keeps the TP16 folded output contract unchanged. +void glmSparseMlaFlashDecodeLatent(th::Tensor const& q, th::Tensor const& kvPool, th::Tensor const& indices, + th::Tensor const& seqLensKv, th::Tensor const& mlaBmm1Scale, th::Tensor const& mlaBmm2Scale, th::Tensor& output, + int64_t seqLenQ, int64_t topK, std::optional const& blockTable, int64_t tokensPerBlock, + int64_t pageIndexScale, int64_t layerOffset, bool gatherBeforeWait, bool triggerAfterWait) +{ + TORCH_CHECK(tensorrt_llm::common::isSM100Family(), "glm_sparse_mla_flash_decode_latent needs an SM100-family GPU"); + TORCH_CHECK(q.is_cuda() && kvPool.is_cuda() && indices.is_cuda() && seqLensKv.is_cuda() && mlaBmm1Scale.is_cuda() + && mlaBmm2Scale.is_cuda() && output.is_cuda(), + "glm_sparse_mla_flash_decode_latent takes CUDA tensors"); + auto const device = q.device(); + TORCH_CHECK(kvPool.device() == device && indices.device() == device && seqLensKv.device() == device + && mlaBmm1Scale.device() == device && mlaBmm2Scale.device() == device && output.device() == device, + "glm_sparse_mla_flash_decode_latent takes tensors on one device"); + + TORCH_CHECK(isE4m3Storage(q), "q must hold E4M3 values (float8_e4m3fn or uint8)"); + TORCH_CHECK(q.dim() == 3 && q.size(1) == tk::kGlmSparseMlaNumHeadsWide && q.size(2) == tk::kGlmSparseMlaHeadDimQk + && q.is_contiguous() && isBulkCopyAligned(q), + "q must be a contiguous, 16-byte aligned [tokens, ", tk::kGlmSparseMlaNumHeadsWide, ", ", + tk::kGlmSparseMlaHeadDimQk, "] tensor"); + int64_t const numTokens = q.size(0); + TORCH_CHECK(numTokens <= tk::kGlmSparseMlaMaxWideTokens, "glm_sparse_mla_flash_decode_latent takes at most ", + tk::kGlmSparseMlaMaxWideTokens, " query tokens, got ", numTokens); + TORCH_CHECK(seqLenQ > 0 && seqLenQ <= kInt32Max && numTokens % seqLenQ == 0, + "seq_len_q must be positive and divide the query tokens; got ", seqLenQ, " for ", numTokens, " tokens"); + int64_t const numSeqs = numTokens / seqLenQ; + + TORCH_CHECK(isE4m3Storage(kvPool), "kv_pool must hold E4M3 values (float8_e4m3fn or uint8)"); + TORCH_CHECK(kvPool.is_contiguous() && isBulkCopyAligned(kvPool) && kvPool.numel() % tk::kGlmSparseMlaHeadDimQk == 0, + "kv_pool must be a contiguous, 16-byte aligned pool of ", tk::kGlmSparseMlaHeadDimQk, "-byte token rows"); + TORCH_CHECK(kvPool.numel() / tk::kGlmSparseMlaHeadDimQk <= kInt32Max, "kv_pool rows must be int32-addressable"); + + TORCH_CHECK(topK > 0 && topK <= tk::kGlmSparseMlaMaxTopK, "top_k must be in [1, ", tk::kGlmSparseMlaMaxTopK, + "], got ", topK); + TORCH_CHECK(indices.scalar_type() == th::kInt32 && indices.dim() == 2 && indices.size(0) == numTokens + && indices.size(1) >= topK && indices.stride(1) == 1 && indices.stride(0) >= topK + && indices.stride(0) <= kInt32Max, + "indices must be an int32 [tokens, >= top_k] tensor with unit column stride"); + TORCH_CHECK(seqLensKv.scalar_type() == th::kInt32 && seqLensKv.dim() == 1 && seqLensKv.size(0) >= numSeqs + && seqLensKv.is_contiguous(), + "seq_lens_kv must be a contiguous int32 tensor with one entry per sequence"); + TORCH_CHECK(mlaBmm1Scale.scalar_type() == th::kFloat32 && mlaBmm1Scale.is_contiguous() && mlaBmm1Scale.numel() >= 2, + "mla_bmm1_scale must hold two contiguous float32 values"); + TORCH_CHECK(mlaBmm2Scale.scalar_type() == th::kFloat32 && mlaBmm2Scale.numel() >= 1, + "mla_bmm2_scale must hold a float32 value"); + + TORCH_CHECK(output.scalar_type() == th::kBFloat16 && output.dim() == 3 && output.size(0) >= numTokens + && output.size(1) == tk::kGlmSparseMlaNumHeadsWide && output.size(2) == tk::kGlmSparseMlaHeadDimV + && output.stride(2) == 1 && output.stride(1) == tk::kGlmSparseMlaHeadDimV + && output.stride(0) >= tk::kGlmSparseMlaNumHeadsWide * tk::kGlmSparseMlaHeadDimV + && output.stride(0) % 2 == 0 && reinterpret_cast(output.data_ptr()) % 4 == 0, + "output must be 4-byte aligned BF16 [>= tokens, 16, 512] with contiguous head/latent axes and even row stride"); + + tk::GlmSparseMlaFlashDecodeParams params{}; + params.q = q.data_ptr(); + params.kvPool = kvPool.data_ptr(); + params.indices = indices.data_ptr(); + params.seqLensKv = seqLensKv.data_ptr(); + params.scaleSoftmaxLog2 = mlaBmm1Scale.data_ptr() + 1; + params.outputScale = mlaBmm2Scale.data_ptr(); + params.vbProj = nullptr; + params.out = output.data_ptr(); + params.outStride = output.stride(0); + params.numTokensQ = static_cast(numTokens); + params.seqLenQ = static_cast(seqLenQ); + params.indicesStride = static_cast(indices.stride(0)); + params.topK = static_cast(topK); + params.gatherBeforeWait = gatherBeforeWait; + params.triggerAfterWait = triggerAfterWait; + if (blockTable.has_value()) + { + auto const& table = blockTable.value(); + TORCH_CHECK(table.is_cuda() && table.device() == device && table.scalar_type() == th::kInt32 && table.dim() == 2 + && table.stride(1) == 1 && table.size(0) >= numSeqs && table.size(1) > 0 && table.size(1) <= kInt32Max + && table.stride(0) >= table.size(1), + "block_table must be an int32 CUDA [sequences, max_blocks] tensor with unit column stride"); + TORCH_CHECK(tokensPerBlock > 0 && pageIndexScale > 0 && layerOffset >= 0, + "invalid block_table geometry: ", "tokens_per_block ", tokensPerBlock, ", page_index_scale ", + pageIndexScale, ", layer_offset ", layerOffset); + int64_t const strideFactor = pageIndexScale * tokensPerBlock; + int64_t const rowOffset = layerOffset * tokensPerBlock; + TORCH_CHECK(strideFactor <= kInt32Max && rowOffset <= kInt32Max, "block_table rows must be int32-addressable"); + params.remapBlockTable = table.data_ptr(); + params.remapBlockTableStride = table.stride(0); + params.remapMaxBlocks = static_cast(table.size(1)); + params.remapTokensPerBlock = static_cast(tokensPerBlock); + params.remapStrideFactor = static_cast(strideFactor); + params.remapLayerOffset = static_cast(rowOffset); + } + else + { + TORCH_CHECK(!gatherBeforeWait, "gather_before_wait needs block_table"); + } + + auto const stream = at::cuda::getCurrentCUDAStream(q.get_device()); + tk::invokeGlmSparseMlaFlashDecodeLatent(params, stream); +} + } // namespace torch_ext TRTLLM_NAMESPACE_END @@ -170,9 +268,15 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) "Tensor mla_bmm1_scale, Tensor mla_bmm2_scale, Tensor v_b_proj, Tensor(a!) output, int seq_len_q, int top_k, " "Tensor? block_table=None, int tokens_per_block=0, int page_index_scale=0, int layer_offset=0, " "bool gather_before_wait=False) -> ()"); + m.def( + "glm_sparse_mla_flash_decode_latent(Tensor q, Tensor kv_pool, Tensor indices, Tensor seq_lens_kv, " + "Tensor mla_bmm1_scale, Tensor mla_bmm2_scale, Tensor(a!) output, int seq_len_q, int top_k, " + "Tensor? block_table=None, int tokens_per_block=0, int page_index_scale=0, int layer_offset=0, " + "bool gather_before_wait=False, bool trigger_after_wait=False) -> ()"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) { m.impl("glm_sparse_mla_flash_decode", &tensorrt_llm::torch_ext::glmSparseMlaFlashDecode); + m.impl("glm_sparse_mla_flash_decode_latent", &tensorrt_llm::torch_ext::glmSparseMlaFlashDecodeLatent); } diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/glmVerifyHeadAcceptOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmVerifyHeadAcceptOp.cpp similarity index 98% rename from cpp/tensorrt_llm/thop/glmMoeDsa/glmVerifyHeadAcceptOp.cpp rename to cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmVerifyHeadAcceptOp.cpp index cfeb6d2ea623..f3479297194f 100644 --- a/cpp/tensorrt_llm/thop/glmMoeDsa/glmVerifyHeadAcceptOp.cpp +++ b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmVerifyHeadAcceptOp.cpp @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/glmMoeDsa/glmVerifyHeadAccept.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmVerifyHeadAccept.h" #include "tensorrt_llm/thop/thUtils.h" #include diff --git a/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmVerifyHeadSampleOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmVerifyHeadSampleOp.cpp new file mode 100644 index 000000000000..707a7641e8c8 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmMoeDsaRubin/glmVerifyHeadSampleOp.cpp @@ -0,0 +1,82 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/glmMoeDsaRubin/glmVerifyHeadSample.h" +#include "tensorrt_llm/thop/thUtils.h" +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ +namespace +{ +void checkArray(at::Tensor const& tensor, at::Device device, at::ScalarType dtype, int64_t minimum, char const* name) +{ + TORCH_CHECK(tensor.device() == device && tensor.scalar_type() == dtype && tensor.is_contiguous() + && tensor.dim() == 1 && tensor.numel() >= minimum, + name, " must be a contiguous CUDA vector of the required dtype and length"); +} +} // namespace + +at::Tensor glmVerifyHeadSampleTokens(at::Tensor const& logits, at::Tensor const& temperatures, at::Tensor const& topK, + at::Tensor const& topP, at::Tensor const& minP, at::Tensor const& seed, at::Tensor const& offset, + at::Tensor& counters, at::Tensor& workspace) +{ + TORCH_CHECK(logits.is_cuda() && logits.scalar_type() == at::kFloat && logits.is_contiguous() && logits.dim() == 2 + && logits.size(0) >= 1 && logits.size(0) <= 24 && logits.size(1) == 154880 + && reinterpret_cast(logits.const_data_ptr()) % 16 == 0, + "GLM verify sampling requires contiguous aligned FP32 logits [1..24, 154880]"); + c10::cuda::CUDAGuard guard(logits.device()); + TORCH_CHECK(tensorrt_llm::common::getSMVersion() == 107, "GLM verify sampling requires SM107"); + int64_t const rows = logits.size(0); + checkArray(temperatures, logits.device(), at::kFloat, rows, "temperatures"); + checkArray(topK, logits.device(), at::kInt, rows, "top_ks"); + checkArray(topP, logits.device(), at::kFloat, rows, "top_ps"); + checkArray(minP, logits.device(), at::kFloat, rows, "min_ps"); + checkArray(seed, logits.device(), at::kLong, 1, "seed"); + checkArray(offset, logits.device(), at::kLong, 1, "offset"); + TORCH_CHECK(seed.numel() == offset.numel() && (seed.numel() == 1 || seed.numel() >= rows), + "RNG arrays must have equal length, either one shared entry or at least one per row"); + checkArray(counters, logits.device(), at::kInt, rows, "row_counters"); + namespace glm = kernels::glm_moe_dsa; + auto const bytes = glm::getGlmVerifyHeadSampleWorkspaceSize(rows, logits.size(1)); + checkArray(workspace, logits.device(), at::kByte, bytes, "workspace"); + TORCH_CHECK(reinterpret_cast(workspace.const_data_ptr()) % 16 == 0, "workspace must be 16-byte aligned"); + auto tokens = at::empty({rows}, logits.options().dtype(at::kInt)); + glm::GlmVerifyHeadSampleParams p{}; + p.logits = logits.const_data_ptr(); + p.temperatures = temperatures.const_data_ptr(); + p.topK = topK.const_data_ptr(); + p.topP = topP.const_data_ptr(); + p.minP = minP.const_data_ptr(); + p.seed = seed.const_data_ptr(); + p.offset = offset.const_data_ptr(); + // The singleton case is per-row, matching current fused sampling's contract. + p.perRowRng = seed.numel() >= rows; + p.rowCounters = counters.mutable_data_ptr(); + p.tokens = tokens.mutable_data_ptr(); + p.workspace = workspace.mutable_data_ptr(); + p.numRows = rows; + p.vocabSize = logits.size(1); + p.stream = at::cuda::getCurrentCUDAStream(logits.get_device()); + glm::invokeGlmVerifyHeadSample(p); + return tokens; +} +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_verify_head_sample_tokens(Tensor logits, Tensor temperatures, Tensor top_ks, Tensor top_ps, " + "Tensor min_ps, Tensor seed, Tensor offset, Tensor(a!) row_counters, Tensor(b!) workspace) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_verify_head_sample_tokens", &tensorrt_llm::torch_ext::glmVerifyHeadSampleTokens); +} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/_full_context.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/_full_context.py new file mode 100644 index 000000000000..0d1f93186c33 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/_full_context.py @@ -0,0 +1,210 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Pre-profile storage and generation indexer for the unsharded Rubin path.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel + +from ._small_batch_draft import DRAFT_SHAPES +from .ops.dsa._full_context_topk import MAX_ROWS, TOPK, full_context_topk, workspace_bytes +from .ops.dsa.fp8_paged_mqa_logits import fp8_paged_mqa_logits +from .ops.dsa.indexer_k_tail import indexer_k_tail +from .ops.dsa.stash_last_accepted_rows import stash_last_accepted_rows +from .ops.gemm._low_m_splitk_fused import q_b_plan +from .ops.gemm.indexer_q_tail import indexer_q_tail + + +class FullContextIndexer: + """One fixed arena reused in stream order by verify and draft attention. + + No cache-manager tensor is retained here: profiling may replace the manager. + ``capacity`` includes the model's conservative speculative-token headroom and + is allocated before memory profiling so the final KV quota accounts for it. + """ + + def __init__(self, decode, capacity: int, device: torch.device) -> None: + self.decode = decode + self.capacity = capacity + self.max_batch = 3 if decode.mapping.tp_size == 4 else 2 + columns = (capacity + 255) // 256 * 256 + self.logits = torch.empty((MAX_ROWS, columns), dtype=torch.float32, device=device) + self.topk = torch.empty((MAX_ROWS, TOPK), dtype=torch.int32, device=device) + self.k_fp8 = torch.empty((MAX_ROWS, 128), dtype=torch.float8_e4m3fn, device=device) + self.k_scale = torch.empty((MAX_ROWS, 1), dtype=torch.float32, device=device) + self.workspace = torch.zeros(workspace_bytes() // 4, dtype=torch.int32, device=device) + + def supports(self, metadata, rows: int) -> bool: + if torch.compiler.is_compiling(): + return False + manager = metadata.kv_cache_manager + batch = metadata.num_seqs + lengths = metadata.gen_indexer_kv_lens_cuda_runtime + table = metadata.indexer_k_cache_block_offsets + shared = metadata.shared_topk_indices + return ( + 1 <= batch <= self.max_batch + and 1 <= rows <= MAX_ROWS + and rows % batch == 0 + and rows == metadata.num_tokens + and ( + rows == batch + or batch <= DRAFT_SHAPES.get((self.decode.mapping.tp_size, rows // batch - 1), 0) + ) + and not metadata.is_ragged_verify + and not metadata.use_expanded_buffers_for_mtp + and not metadata.expand_for_dsl + and manager.tokens_per_block == 64 + and not manager.use_fp4 + and 1 <= manager.max_seq_len <= self.capacity + and 1 <= metadata.get_indexer_max_seq_len() <= self.capacity + and lengths is not None + and lengths.shape == (batch,) + and lengths.dtype == torch.int32 + and lengths.is_contiguous() + and lengths.device == self.logits.device + and table is not None + and table.ndim == 2 + and table.shape[0] >= batch + and table.dtype == torch.int32 + and table.stride(1) == 1 + and table.device == self.logits.device + and shared is not None + and shared.shape[0] >= rows + and shared.shape[1] == TOPK + and metadata.scheduler_metadata_buffer is not None + ) + + def reuse_draft_selection(self, indexer, metadata, draft_state) -> bool: + return ( + draft_state is not None + and indexer is not None + and indexer.mtp_index_share + and metadata.in_mtp_draft_loop + and metadata.indexer_skip_topk + and not metadata.skip_indexer_for_gen_reqs + and metadata.shared_topk_indices is not None + ) + + def select(self, layer, kv_a, q_a, position_ids, metadata, draft_state=None, *, beside_topk): + """Write keys, score the full context, and preserve the current layer's local TopK.""" + decode = self.decode + mla = layer.self_attn + indexer = mla.mqa.indexer + rows, batch = metadata.num_tokens, metadata.num_seqs + next_n = rows // batch + manager = metadata.kv_cache_manager + rope = indexer.rotary_emb + cos_sin = rope.rotary_cos_sin.view(rope.max_positions, -1) + positions = position_ids.reshape(-1)[:rows] + latent_end = mla.q_lora_rank + mla.kv_lora_rank + mla.qk_rope_head_dim + index_k = kv_a[:, latent_end : latent_end + indexer.head_dim] + index_weights = kv_a[:, latent_end + indexer.head_dim :] + k_cache = manager.get_indexer_k_cache_buffers(indexer.layer_idx) + decode._k_tail_fork.record() + with torch.cuda.stream(decode._k_tail_stream): + decode._k_tail_stream.wait_event(decode._k_tail_fork) + indexer_k_tail( + index_k, + positions, + cos_sin, + indexer.k_norm.weight, + indexer.k_norm.bias, + indexer.k_norm.variance_epsilon, + indexer.rope_dim, + not rope.is_neox, + indexer.scale_fmt == "ue8m0", + self.k_fp8[:rows], + self.k_scale[:rows], + k_cache, + metadata.slot_mapping_fp8[:rows], + metadata.slot_mapping_scale[:rows], + rows, + True, + True, + False, + ) + decode._k_tail_done.record() + topk = self.topk[:rows] + if metadata.skip_indexer_for_gen_reqs: + # Short-context bypass still appends K for a later step that runs the indexer. + torch.cuda.current_stream().wait_event(decode._k_tail_done) + topk.copy_(metadata.topk_indices_buffer[:rows]) + beside_topk() + else: + q = q_a.new_empty((rows, indexer.n_heads, indexer.head_dim), dtype=torch.float8_e4m3fn) + q_scale = q_a.new_empty((rows, indexer.n_heads, 1), dtype=torch.float32) + weights = q_a.new_empty((rows, indexer.n_heads), dtype=torch.float32) + indexer_q_tail( + q_a, + decode._wq_b_folded[layer.layer_idx], + positions, + cos_sin, + index_weights, + q, + q_scale, + weights, + mla.q_a_layernorm.variance_epsilon, + indexer.n_heads, + indexer.head_dim, + indexer.rope_dim, + not rope.is_neox, + indexer.scale_fmt == "ue8m0", + indexer.weight_scale_factor, + ) + torch.cuda.current_stream().wait_event(decode._k_tail_done) + capacity = metadata.get_indexer_max_seq_len() + logits = fp8_paged_mqa_logits( + q.view(batch, next_n, indexer.n_heads, indexer.head_dim), + k_cache, + weights, + metadata.gen_indexer_kv_lens_cuda_runtime, + metadata.indexer_k_cache_block_offsets[:batch], + metadata.scheduler_metadata_buffer, + self.logits[:rows], + capacity, + next_n % 2 == 0, + ) + tactic = q_b_plan( + "beside_topk", rows, mla.num_heads_tp * 576, mla.q_lora_rank, q_a.device.index + ) + reserve = ( + (mla.num_heads_tp * 576 + tactic.mma_m - 1) + // tactic.mma_m + * ((rows + tactic.mma_n - 1) // tactic.mma_n) + * tactic.split_k + ) + + def topk_select(): + full_context_topk( + logits, + metadata.gen_indexer_kv_lens_cuda_runtime, + topk, + self.workspace, + next_n=next_n, + max_seq_len=capacity, + sm_reserve=reserve, + deterministic=indexer.top_k.gvr_deterministic, + ) + + maybe_execute_in_parallel( + topk_select, beside_topk, mla.ln_events[0], mla.ln_events[1], mla.aux_stream + ) + + if metadata.in_mtp_draft_loop and indexer.mtp_index_share: + if draft_state is None: + raise RuntimeError("owned MTP attention needs its accepted-row stash") + stash = draft_state.topk_stash[:batch] + accepted = metadata.mtp_num_accepted + if accepted is None: + stash.copy_(topk[next_n - 1 :: next_n]) + else: + stash_last_accepted_rows(topk, accepted[:batch], stash, next_n) + # Only B accepted rows belong to the MTP shared buffer. Never overwrite + # them with all B*next_n rows after compaction. + metadata.shared_topk_indices[:batch].copy_(stash) + else: + metadata.shared_topk_indices[:rows].copy_(topk) + return topk diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/_model.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/_model.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/_model.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/_model.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/_mtp.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/_mtp.py similarity index 95% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/_mtp.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/_mtp.py index e951cfa87c99..d80473b91083 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/_mtp.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/_mtp.py @@ -180,8 +180,26 @@ def forward( attn_metadata: AttentionMetadata, all_rank_num_tokens: Optional[List[int]] = None, spec_metadata: Optional[SpecMetadata] = None, + owned_draft=None, + eh_proj_input: Optional[torch.Tensor] = None, + eh_proj_output: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + row_select=None, **kwargs, ) -> torch.Tensor: + if owned_draft is not None: + hidden_states = owned_draft.forward_step( + input_ids, + position_ids, + hidden_states, + attn_metadata, + eh_proj_input=eh_proj_input, + eh_proj_output=eh_proj_output, + row_select=row_select, + ) + if spec_metadata is not None: + spec_metadata.maybe_capture_hidden_states(0, hidden_states, None) + return hidden_states + def norm_embeds(): return self.enorm(embed_tokens(input_ids)) # emdedding diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/_mtp_worker.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/_mtp_worker.py new file mode 100644 index 000000000000..cac11254463e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/_mtp_worker.py @@ -0,0 +1,568 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GLM-owned MTP worker and optimized draft-loop dispatch.""" + +from __future__ import annotations + +import contextlib +from typing import TYPE_CHECKING, Optional, Tuple + +import torch + +from tensorrt_llm._torch.attention.backends.sparse.params import MTPIndexShareMetadata +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.speculative import MTPEagleWorker +from tensorrt_llm._torch.speculative.eagle3 import _reset_mtp_index_share +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, +) +from tensorrt_llm.llmapi.llm_args import AdvancedSamplingMode + +from ._small_batch_draft import DRAFT_SHAPES, FirstDraftRowSelect +from .ops.spec.verify_head_accept import glm_verify_head_accept + +if TYPE_CHECKING: + from .modeling import DecodeTarget + + +class MTPDraftWorker(MTPEagleWorker): + """The one-model MTP-Eagle worker, preserving main's draft metadata lifecycle. + + The TP16 single-request inline path runs first; eligible TP4/TP8 small batches + then use owned draft kernels. Other steps use the built-in loop. + + The loop keeps the built-in loop's attention-metadata handling: the MTP index-share flags, + the switch to one row per request after draft step 0, and the cleanup. The draft-step + kernels advance the sequence and KV lengths in place, and the draft tokens and the + next-new-tokens rows come from them. + """ + + def __init__( + self, + spec_config, + model_config: ModelConfig, + use_separate_draft_kv_cache: bool, + *, + mapping, + decode_target, + ) -> None: + super().__init__(spec_config, model_config, use_separate_draft_kv_cache, mapping=mapping) + # The core's per-step lookup of the decode target (a bound method rather than the + # core itself, which belongs to the shell and is no submodule of this worker). + self._decode_target = decode_target + # The draft tokens of the last draft loop on the decode target and the next-new-tokens + # rows its last step wrote, for _prepare_next_new_tokens. + self._drafted_rows: Optional[Tuple[torch.Tensor, torch.Tensor]] = None + self._verify_sampler = None + + def prepare_1st_drafter_inputs( + self, + input_ids, + position_ids, + hidden_states, + accepted_tokens, + attn_metadata, + spec_metadata, + draft_model, + ): + """Without context requests draft step 0's input ids are the accepted-token rows as + they are, flattened, rather than a copy behind an empty context part.""" + if not self.is_mtp_eagle or attn_metadata.num_contexts > 0: + return super().prepare_1st_drafter_inputs( + input_ids, + position_ids, + hidden_states, + accepted_tokens, + attn_metadata, + spec_metadata, + draft_model, + ) + return { + "input_ids": accepted_tokens[:, : spec_metadata.runtime_draft_len + 1].flatten(), + "position_ids": position_ids, + "hidden_states": hidden_states, + "attn_metadata": attn_metadata, + "spec_metadata": spec_metadata, + } + + def _prepare_next_new_tokens( + self, + accepted_tokens, + next_draft_tokens, + batch_indices_cuda, + batch_size, + num_accepted_tokens, + ): + """The rows the last draft-step kernel wrote when these draft tokens are that loop's.""" + drafted, self._drafted_rows = self._drafted_rows, None + if drafted is not None and drafted[0] is next_draft_tokens: + return drafted[1] + return super()._prepare_next_new_tokens( + accepted_tokens, + next_draft_tokens, + batch_indices_cuda, + batch_size, + num_accepted_tokens, + ) + + def _sample_tokens_for_batch(self, logits, spec_metadata, num_contexts, batch_size): + """Sample gathered generation logits with the owned FULL sampler when eligible.""" + sampler = self._verify_sampler + if ( + sampler is not None + and num_contexts == 0 + and spec_metadata is not None + and not spec_metadata.is_ragged_verify + and not spec_metadata.is_all_greedy_sample + and spec_metadata.advanced_sampling_mode == AdvancedSamplingMode.FULL + and isinstance(logits, torch.Tensor) + and not torch.compiler.is_compiling() + ): + rows = logits.shape[0] + parameters = ( + spec_metadata.temperatures[:rows], + spec_metadata.top_ks[:rows], + spec_metadata.top_ps[:rows], + spec_metadata.min_ps[:rows], + *self._rng_state_per_token(spec_metadata, rows), + ) + if sampler.supported(logits, *parameters): + return sampler.sample(logits, *parameters) + return super()._sample_tokens_for_batch(logits, spec_metadata, num_contexts, batch_size) + + def _sample_and_accept_draft_tokens_base( + self, + logits, + draft_tokens, + num_contexts, + batch_size, + spec_metadata, + ): + """Strict acceptance and the forced accepted-token count in one kernel + (``glm_verify_head_accept``) for a step of generation requests only, with the worker's + target-token sampling; ragged verification and context requests take the built-in + chain. The forced count follows ``_apply_force_accepted_tokens``: none during eager + CUDA-graph warm-up, the integer part plus one, and with a fractional part the same pool + draw and counter step.""" + if spec_metadata is None or spec_metadata.is_ragged_verify or num_contexts != 0: + return super()._sample_and_accept_draft_tokens_base( + logits, draft_tokens, num_contexts, batch_size, spec_metadata + ) + if logits.dim() == 1: + logits = logits.unsqueeze(0) + tokens_per_request = draft_tokens.shape[-1] + 1 + target_tokens = self._sample_tokens_for_batch(logits, spec_metadata, 0, batch_size) + mode, base, frac = 0, 0, 0.0 + warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() + if self.force_num_accepted_tokens != 0.0 and not warmup: + int_part = int(self.force_num_accepted_tokens) + frac = self.force_num_accepted_tokens - int_part + base = min(int_part + 1, tokens_per_request) + mode = 2 if frac > 0.0 and base < tokens_per_request else 1 + pool = counter = None + if mode == 2: + self._ensure_force_accept_rng_state(logits.device) + pool, counter = self._force_accept_rng_pool, self._force_accept_rng_counter + accepted_tokens = torch.empty( + (batch_size, tokens_per_request), dtype=torch.int, device=logits.device + ) + num_accepted_tokens = glm_verify_head_accept( + target_tokens, + draft_tokens, + 0, + accepted_tokens, + pool, + counter, + mode, + base, + tokens_per_request, + frac, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + ) + return accepted_tokens, num_accepted_tokens + + def _forward_linear_draft_loop( + self, + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + accepted_tokens=None, + pre_draft_seq_lens=None, + ): + self._drafted_rows = None + decode = self._drafting_target( + inputs, + attn_metadata, + spec_metadata, + draft_kv_cache_manager, + num_contexts, + batch_size, + num_accepted_tokens, + ) + if decode is not None: + drafted = self._draft_on_target( + decode, + inputs, + attn_metadata, + spec_metadata, + draft_kv_cache_manager, + batch_size, + num_accepted_tokens, + ) + return drafted, None + helper = self._small_batch_target( + inputs, + attn_metadata, + spec_metadata, + draft_model, + num_contexts, + batch_size, + num_accepted_tokens, + ) + if helper is not None: + drafted = self._forward_small_batch_draft_loop( + helper, + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + batch_size, + num_accepted_tokens, + ) + if drafted is not None: + return drafted, None + return super()._forward_linear_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + accepted_tokens, + pre_draft_seq_lens, + ) + + def _small_batch_target( + self, + inputs, + attn_metadata, + spec_metadata, + draft_model, + num_contexts, + batch_size, + num_accepted_tokens, + ): + """Select the small-batch bridge before any draft layer or collective.""" + draft_len = spec_metadata.runtime_draft_len + rows = batch_size * (draft_len + 1) + ids, positions, hidden = ( + inputs["input_ids"], + inputs["position_ids"], + inputs["hidden_states"], + ) + batch_indices = spec_metadata.batch_indices_cuda + max_batch = DRAFT_SHAPES.get((self.mapping.tp_size, self.max_draft_len), 0) + if not ( + self.is_mtp_eagle + and num_contexts == 0 + and 1 <= batch_size <= max_batch + and draft_len == self.max_draft_len + and getattr(self.spec_config, "draft_len_schedule", None) is None + and self.guided_decoder is None + and self.sa_enhancer is None + and self._d2t is None + and not spec_metadata.wants_advanced_draft_sampling + and not spec_metadata.is_ragged_verify + and not torch.compiler.is_compiling() + and isinstance(attn_metadata, MTPIndexShareMetadata) + and attn_metadata.num_tokens == rows + and ids.dtype == torch.int32 + and tuple(ids.shape) == (rows,) + and ids.is_cuda + and ids.is_contiguous() + and positions.dtype in (torch.int32, torch.int64) + and tuple(positions.shape) == (rows,) + and positions.is_contiguous() + and positions.device == ids.device + and hidden.dtype == torch.bfloat16 + and tuple(hidden.shape) == (rows, 6144) + and hidden.device == ids.device + and hidden.is_contiguous() + and hidden.data_ptr() % 32 == 0 + and num_accepted_tokens.dtype == torch.int32 + and num_accepted_tokens.dim() == 1 + and num_accepted_tokens.numel() >= batch_size + and num_accepted_tokens.device == ids.device + and num_accepted_tokens.is_contiguous() + and batch_indices.dtype == torch.int32 + and batch_indices.dim() == 1 + and batch_indices.numel() >= batch_size + and batch_indices.device == ids.device + and batch_indices.is_contiguous() + ): + return None + decode = self._decode_target(attn_metadata) + helper = getattr(decode, "_small_batch_draft", None) + if ( + helper is None + or helper.DRAFT_LEN != draft_len + or batch_size > helper.MAX_BATCH + or len(draft_model.mtp_layers) != 1 + or helper.layer is not draft_model.mtp_layers[0] + or helper.lm_head is not draft_model.lm_head + ): + return None + return helper + + def _forward_small_batch_draft_loop( + self, + helper, + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + batch_size, + num_accepted_tokens, + ): + """Keep main's metadata lifecycle around owned full-attention draft steps. + + None means only the preflight declined: no draft forward has executed. + Once the first layer starts, failures propagate without replaying it. + """ + from tensorrt_llm._torch.attention.backends.sparse.dsa import is_dsa_cache_manager + + draft_len = spec_metadata.runtime_draft_len + first_rows = batch_size * (draft_len + 1) + accepted_tokens = inputs["input_ids"].view(batch_size, draft_len + 1) + batch_indices = spec_metadata.batch_indices_cuda[:batch_size] + num_accepted = num_accepted_tokens[:batch_size] + selector = FirstDraftRowSelect(batch_indices, num_accepted, draft_len + 1) + # Do not modify the caller's inputs if the rebound-manager preflight declines. + inputs = dict(inputs) + with contextlib.ExitStack() as scope: + attn_metadata.set_in_mtp_draft_loop(True) + attn_metadata.set_mtp_num_accepted(num_accepted_tokens) + scope.callback(_reset_mtp_index_share, attn_metadata) + attn_metadata = scope.enter_context( + self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager) + ) + # A V2 draft view can differ from the target manager without a + # separately configured draft cache. Validate the actual rebound views. + if not helper.supports(attn_metadata, first_rows, batch_size): + return None + if is_dsa_cache_manager(draft_kv_cache_manager): + attn_metadata.on_update_kv_lens() + inputs["attn_metadata"] = attn_metadata + inputs["owned_draft"] = helper + inputs["row_select"] = selector + draft_tokens = accepted_tokens.new_empty((batch_size, draft_len)) + next_new_tokens = None + for step in range(draft_len): + attn_metadata.set_skip_topk(step > 0) + hidden, _ = self._run_draft_forward(draft_model, inputs, spec_metadata, step) + # The layer already selected both hidden and residual BEFORE MoE. + # A second gather here would select the wrong request (or go OOB). + if step == 0 and selector.gather_ids is None: + raise RuntimeError("small-batch MTP did not select its first-step rows") + logits, hidden_half = helper.head(hidden, next_step=step + 1 < draft_len) + if step == 0: + attn_metadata._seq_lens[:batch_size].fill_(1) + attn_metadata.on_update() + attn_metadata.num_contexts = 0 + next_projection, next_positions, next_new_tokens = helper.sample_step( + logits, + hidden, + draft_tokens, + step, + inputs["position_ids"], + selector.gather_ids if step == 0 else None, + attn_metadata.seq_lens_cuda[:batch_size], + attn_metadata.kv_lens_cuda[:batch_size], + num_accepted, + accepted_tokens, + batch_indices, + hidden_half=hidden_half, + ) + # The fused op just wrote seq/KV lengths. Main owns derived DSA + # slots, schedules and dense selection; update even the final step. + attn_metadata.update_for_spec_dec() + if step == 0: + attn_metadata.use_spec_decoding = False + if step + 1 < draft_len: + inputs = { + "input_ids": draft_tokens[:, step], + "position_ids": next_positions, + "hidden_states": hidden, + "attn_metadata": attn_metadata, + "spec_metadata": spec_metadata, + "owned_draft": helper, + "eh_proj_output" if helper.use_table else "eh_proj_input": next_projection, + } + self._drafted_rows = (draft_tokens, next_new_tokens) + return draft_tokens + + def _drafting_target( + self, + inputs, + attn_metadata, + spec_metadata, + draft_kv_cache_manager, + num_contexts, + batch_size, + num_accepted_tokens, + ) -> Optional[DecodeTarget]: + """The decode target when its kernels draft this step, else None. + + They draft a generation step of one request whose draft layer shares the target's KV + pool, through the configured draft length with no draft-length schedule (the draft + steps after the first skip the MTP layer's indexer K write: the next step's draft + step 0 writes those positions before it reads them), greedily (argmax drafts), with + nothing that rewrites the drafts (guided decoding, SA, a draft vocabulary map), and + when the decode target's kernels serve the rows. + """ + draft_len = spec_metadata.runtime_draft_len + rows = batch_size * (draft_len + 1) + input_ids = inputs["input_ids"] + position_ids = inputs["position_ids"] + if not ( + self.is_mtp_eagle + and num_contexts == 0 + and batch_size == 1 + and draft_kv_cache_manager is None + and draft_len == self.max_draft_len + and getattr(self.spec_config, "draft_len_schedule", None) is None + and self.guided_decoder is None + and self.sa_enhancer is None + and self._d2t is None + and not spec_metadata.wants_advanced_draft_sampling + and not spec_metadata.is_ragged_verify + and attn_metadata.num_tokens == rows + and inputs["hidden_states"].shape[0] == rows + and input_ids.dtype == torch.int32 + and tuple(input_ids.shape) == (rows,) + and input_ids.is_contiguous() + and position_ids.dtype in (torch.int32, torch.int64) + and tuple(position_ids.shape) == (rows,) + and position_ids.is_contiguous() + and num_accepted_tokens.dtype == torch.int32 + and num_accepted_tokens.is_contiguous() + and spec_metadata.batch_indices_cuda.dtype == torch.int32 + ): + return None + decode = self._decode_target(attn_metadata) + if decode is None or not decode.drafts(rows): + return None + return decode + + def _draft_on_target( + self, + decode: DecodeTarget, + inputs, + attn_metadata, + spec_metadata, + draft_kv_cache_manager, + batch_size, + num_accepted_tokens, + ) -> torch.Tensor: + """The draft loop on the decode target's kernels: the draft tokens [batch, draft_len]. + + Each step's last kernel writes the next step's sequence and KV lengths in place. The + steps after the first run no indexer and read nothing else of the DSA decode metadata, + so the built-in loop's per-step rebuild of it (``update_for_spec_dec``) is left out; + with a context short enough that the metadata skips the indexer, only the next step's + dense selection is refreshed from the new lengths. The host side of draft step 0's + switch to one row per request comes before that. + """ + draft_len = spec_metadata.runtime_draft_len + # The accepted tokens' rows: draft step 0's input ids (prepare_1st_drafter_inputs). + accepted_tokens = inputs["input_ids"].view(batch_size, draft_len + 1) + batch_indices = spec_metadata.batch_indices_cuda[:batch_size] + num_accepted = num_accepted_tokens[:batch_size] + draft_tokens = accepted_tokens.new_empty((batch_size, draft_len)) + positions = inputs["position_ids"] + eh_proj_output = gather_ids = pool_rows = next_new_tokens = None + uses_mtp_index_share = isinstance(attn_metadata, MTPIndexShareMetadata) + with contextlib.ExitStack() as draft_scope: + if uses_mtp_index_share: + attn_metadata.set_in_mtp_draft_loop(True) + attn_metadata.set_mtp_num_accepted(num_accepted_tokens) + draft_scope.callback(_reset_mtp_index_share, attn_metadata) + attn_metadata = draft_scope.enter_context( + self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager) + ) + # The draft loop's copy of the sequence lengths (prepare_for_spec_dec) and the KV + # lengths, which the draft-step kernels advance. + seq_lens = attn_metadata.seq_lens_cuda[:batch_size] + kv_lens = attn_metadata.kv_lens_cuda[:batch_size] + for i in range(draft_len): + if uses_mtp_index_share: + attn_metadata.set_skip_topk(i > 0) + if i == 0: + hidden_states, gather_ids, pool_rows = decode.draft_first_step( + attn_metadata, + inputs["input_ids"], + positions, + inputs["hidden_states"], + batch_indices, + num_accepted, + draft_len, + ) + else: + hidden_states = decode.draft_next_step( + attn_metadata, eh_proj_output, pool_rows, join_selection=i == 1 + ) + last = i == draft_len - 1 + logits, hidden_half = decode.draft_head(hidden_states, next_step=not last) + rewind_accepted = num_accepted if i == 0 else None + if i == 0: + # One row per request from here on, on the host (no context requests to + # turn into generation ones); the kernel below fills seq_lens. + attn_metadata._seq_lens[:batch_size].fill_(1) + attn_metadata.on_update() + if last: + next_new_tokens = decode.draft_sample_glue( + logits, + draft_tokens, + i, + seq_lens, + kv_lens, + rewind_accepted, + accepted_tokens, + batch_indices, + num_accepted, + ) + else: + eh_proj_output, positions = decode.draft_ehproj_table( + logits, + hidden_half, + draft_tokens, + i, + positions, + gather_ids if i == 0 else None, + seq_lens, + kv_lens, + rewind_accepted, + ) + if not last and attn_metadata.skip_indexer_for_gen_reqs: + attn_metadata.prepare_dense_topk_indices( + attn_metadata.kv_lens_cuda, device=True + ) + if i == 0: + attn_metadata.use_spec_decoding = False + self._drafted_rows = (draft_tokens, next_new_tokens) + return draft_tokens diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/_small_batch_draft.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/_small_batch_draft.py new file mode 100644 index 000000000000..65ef484da047 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/_small_batch_draft.py @@ -0,0 +1,475 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Owned TP4/TP8 MTP kernels and their pre-profile state. + +This is a plain Python object, not an nn.Module: its callbacks may reference the +target without registering a parent module under the draft layer. The target +owns the full-context attention arena; this helper owns only draft operands and +small buffers. All eligibility decisions precede the first draft KV write. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from tensorrt_llm._torch.distributed.ops import _get_mnnvl_workspace_comm +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.modules.linear import TensorParallelMode +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.utils import EventType + +from .ops.comm._mnnvl_workspace import hop_workspace +from .ops.comm.mnnvl_moe_finalize_hop import mnnvl_moe_finalize_hop +from .ops.gemm.draft_hidden_half import draft_hidden_half +from .ops.gemm.low_m_gemm import low_m_gemm +from .ops.gemm.router_gemm import router_gemm +from .ops.moe.lowm_bf16_moe import glm_lowm_bf16_moe +from .ops.spec._draft_exchange import DraftExchangeWorkspace, _all_succeeded +from .ops.spec._draft_token_table import build_draft_token_table +from .ops.spec.draft_ehproj_table import glm_draft_ehproj_table +from .ops.spec.draft_glue_rmsnorm_cat import glm_draft_glue_rmsnorm_cat +from .ops.spec.draft_sample_glue import glm_draft_sample_glue +from .ops.spec.gather_last_accepted_rows import glm_gather_last_accepted_rows + +if TYPE_CHECKING: + from .modeling import DecodeTarget + from .weights import DraftWeights + +DRAFT_SHAPES = {(4, 5): 3, (8, 3): 2, (8, 7): 1} # (TP, draft length): maximum batch + + +class FirstDraftRowSelect: + """Select hidden and residual together after the complete first attention.""" + + def __init__(self, batch_indices, num_accepted, tokens_per_request: int) -> None: + self.batch_indices = batch_indices + self.num_accepted = num_accepted + self.tokens_per_request = tokens_per_request + self.gather_ids = None + + def __call__(self, hidden, residual): + if self.gather_ids is not None: + raise RuntimeError("the first draft rows were already selected") + self.gather_ids, hidden, residual = glm_gather_last_accepted_rows( + hidden, + self.batch_indices, + self.num_accepted, + self.tokens_per_request, + 0, + second=residual, + ) + return hidden, residual + + +def _bf16_matrix(tensor, shape) -> bool: + return ( + tensor.dtype == torch.bfloat16 + and tuple(tensor.shape) == shape + and tensor.is_cuda + and tensor.is_contiguous() + and tensor.data_ptr() % 32 == 0 + ) + + +def _ordinary_norm(norm, hidden: int) -> bool: + return ( + _bf16_matrix(norm.weight, (hidden,)) + and not norm.use_gemma + and not norm.use_cuda_tile + and not norm.return_hp_output + and not norm.is_nvfp4 + ) + + +def _static_layout_supported(decode, layer, lm_head, weights) -> bool: + mapping = decode.mapping + config = decode.core.model_config.pretrained_config + hidden, experts = 6144, 256 + inter = 2048 // mapping.tp_size + columns = 2 * hidden // mapping.tp_size + embed = decode.core.embed_tokens.weight + gate = layer.mlp.gate + bias = gate.e_score_correction_bias + norms = ( + layer.enorm, + layer.hnorm, + layer.input_layernorm, + layer.post_attention_layernorm, + layer.shared_head.norm, + ) + indexer = layer.self_attn.mqa.indexer + return ( + mapping.tp_size in (4, 8) + and mapping.tp_size == mapping.world_size + and mapping.moe_ep_size == 1 + and mapping.pp_size == mapping.cp_size == 1 + and not mapping.enable_attention_dp + and config.hidden_size == hidden + and config.n_routed_experts == experts + and config.n_group == 1 + and config.num_experts_per_tok == 8 + and config.moe_intermediate_size // mapping.tp_size == inter + and indexer is not None + and indexer.mtp_index_share + and all(_ordinary_norm(norm, hidden) for norm in norms) + and layer.enorm.variance_epsilon == layer.hnorm.variance_epsilon + and layer.eh_proj.tp_mode == TensorParallelMode.ROW + and layer.eh_proj.reduce_output + and layer.eh_proj.bias is None + and _bf16_matrix(layer.eh_proj.weight, (hidden, columns)) + and (mapping.tp_size == 4 or _bf16_matrix(weights.eh_proj, (hidden, 2 * hidden))) + and _bf16_matrix(weights.w3_w1, (experts, 2 * inter, hidden)) + and _bf16_matrix(weights.w2, (experts, hidden, inter)) + and _bf16_matrix(gate.weight, (experts, hidden)) + and bias.dtype in (torch.float32, torch.bfloat16) + and bias.is_cuda + and bias.is_contiguous() + and tuple(bias.shape) == (experts,) + and lm_head.tp_mode == TensorParallelMode.COLUMN + and getattr(lm_head, "padding_size", 1) == 0 + and lm_head.weight.shape[0] % 8 == 0 + and lm_head.weight.shape[0] * mapping.tp_size == lm_head.num_embeddings + and _bf16_matrix(lm_head.weight, (lm_head.num_embeddings // mapping.tp_size, hidden)) + and _bf16_matrix(embed, (lm_head.num_embeddings, hidden)) + ) + + +class SmallBatchDraftKernels: + """TP4 B<=3/K5 or TP8 B<=2/K3 and B1/K7 state, allocated before profiling.""" + + MAX_ROWS = 24 + + def __init__( + self, + decode, + layer, + lm_head, + weights, + sample_exchange, + table_exchange=None, + *, + draft_len: int, + ) -> None: + self.decode = decode + self.layer = layer + self.lm_head = lm_head + self.sample_exchange = sample_exchange + self.table_exchange = table_exchange + self.use_table = table_exchange is not None + self.DRAFT_LEN = draft_len + self.MAX_BATCH = DRAFT_SHAPES[decode.mapping.tp_size, draft_len] + self.embed_table = decode.core.embed_tokens.weight.detach() + self.enorm_weight = layer.enorm.weight.detach() + self.hnorm_weight = layer.hnorm.weight.detach() + self.w3_w1, self.w2 = weights.w3_w1, weights.w2 + config = decode.core.model_config.pretrained_config + self.top_k = config.num_experts_per_tok + self.routed_scaling_factor = config.routed_scaling_factor + self.hidden_size = config.hidden_size + columns = 2 * self.hidden_size // decode.mapping.tp_size + self.window = (decode.mapping.tp_rank * columns, (decode.mapping.tp_rank + 1) * columns) + self.pdl = get_env_enable_pdl() + self.zero_residual = torch.zeros( + (self.MAX_ROWS, self.hidden_size), dtype=torch.bfloat16, device=self.embed_table.device + ) + # Never alias the target's full-row TopK output during accepted-row compaction. + self.topk_stash = torch.empty( + (self.MAX_BATCH, config.index_topk), dtype=torch.int32, device=self.embed_table.device + ) + if self.use_table: + self.token_table, self.hidden_rows = build_draft_token_table( + weights.eh_proj, + self.embed_table, + self.enorm_weight, + layer.enorm.variance_epsilon, + decode.mapping.tp_rank, + decode.mapping.tp_size, + lm_head.weight.shape[0], + ) + self.half_start = torch.cuda.Event() + self.half_done = torch.cuda.Event() + self.identity_routing = {} + for rows in range(1, self.MAX_BATCH + 1): + expert_weights = torch.zeros((rows, self.top_k), dtype=torch.bfloat16) + expert_rows = torch.full((rows, self.top_k), -1, dtype=torch.int32) + expert_weights[:, 0] = 1 + expert_rows[:, 0] = torch.arange(rows, dtype=torch.int32) + self.identity_routing[rows] = ( + expert_weights.to(self.embed_table.device), + expert_rows.to(self.embed_table.device), + ) + mla = layer.self_attn + self.l2_after_fused_a = [ + tensor + for tensor in (decode._q_b_absorbed[layer.layer_idx], mla.v_b_proj, mla.o_proj.weight) + if tensor.is_contiguous() + ] + + def supports(self, metadata, first_rows: int, batch_size: int) -> bool: + return ( + 1 <= batch_size <= self.MAX_BATCH + and first_rows == batch_size * (self.DRAFT_LEN + 1) + and first_rows <= self.MAX_ROWS + and self.decode.supports_small_batch_draft(metadata, first_rows, batch_size, self.layer) + ) + + def forward_step( + self, + input_ids, + position_ids, + hidden_states, + attn_metadata, + *, + eh_proj_input=None, + eh_proj_output=None, + row_select=None, + ): + """Full first attention, optional row selection, then the B-row BF16 MoE.""" + if eh_proj_input is not None and eh_proj_output is not None: + raise ValueError("eh_proj_input and eh_proj_output are mutually exclusive") + layer = self.layer + rows = hidden_states.shape[0] + if eh_proj_output is None: + if eh_proj_input is None: + eh_proj_input = glm_draft_glue_rmsnorm_cat( + self.embed_table, + input_ids, + hidden_states, + self.enorm_weight, + self.hnorm_weight, + layer.enorm.variance_epsilon, + *self.window, + self.pdl, + ) + partial = low_m_gemm( + eh_proj_input, + layer.eh_proj.weight, + eh_proj_input.new_empty((rows, self.hidden_size)), + ) + hidden_states, residual = self.decode._hop( + layer.input_layernorm, partial, self.zero_residual[:rows] + ) + else: + hidden_states, residual = eh_proj_output + attn_out = self.decode._attention( + layer, position_ids, hidden_states, attn_metadata, False, draft_state=self + ) + hidden_states, residual = self.decode._post_attention_hop(layer, attn_out, residual) + if row_select is not None: + hidden_states, residual = row_select(hidden_states, residual) + return self._moe_postnorm(hidden_states, residual) + + def _moe_postnorm(self, hidden_states, residual): + mlp = self.layer.mlp + gate = mlp.gate + + def routed(): + logits = router_gemm(hidden_states, gate.weight) + return glm_lowm_bf16_moe( + hidden_states, + logits, + gate.e_score_correction_bias, + self.w3_w1, + self.w2, + self.top_k, + self.routed_scaling_factor, + ) + + routed_output, shared_output = maybe_execute_in_parallel( + routed, + lambda: self.decode._shared_expert(mlp, hidden_states), + mlp.event_dict[EventType.Main], + mlp.event_dict[EventType.MoeShared], + mlp.aux_stream, + ) + if self.use_table: + rows = hidden_states.shape[0] + expert_weights, expert_rows = self.identity_routing[rows] + comm_buffer, buffer_flags = hop_workspace(self.decode.mapping, rows, self.hidden_size) + norm = self.layer.shared_head.norm + normed, _ = mnnvl_moe_finalize_hop( + routed_output, + expert_weights, + expert_rows, + shared_output, + residual, + norm.weight, + norm.variance_epsilon, + comm_buffer, + buffer_flags, + ) + return normed + # Match the plain TP4 path's BF16 local sum, then FP32 rank reduction. + # Fake identity-routing finalize would introduce different rank rounding. + partial = routed_output + shared_output + normed, _ = self.decode._hop(self.layer.shared_head.norm, partial, residual) + return normed + + def head(self, hidden_states, *, next_step=False): + # The layer already applied shared_head.norm; do not normalize again. + if not (self.use_table and next_step): + return self.decode._lm_head_shard(self.lm_head, hidden_states), None + rows = hidden_states.shape[0] + hidden_half = torch.empty( + (rows, self.hidden_rows.shape[0]), dtype=torch.float32, device=hidden_states.device + ) + self.half_start.record() + # B2 launches its LM head first; B1 folds hnorm into the hidden-half GEMM. + logits = self.decode._lm_head_shard(self.lm_head, hidden_states) if rows > 1 else None + aux_stream = self.layer.mlp.aux_stream + with torch.cuda.stream(aux_stream): + aux_stream.wait_event(self.half_start) + normed = hidden_states if rows == 1 else self.layer.hnorm(hidden_states) + draft_hidden_half( + normed, + self.hidden_rows, + hidden_half, + self.hnorm_weight if rows == 1 else None, + self.layer.hnorm.variance_epsilon if rows == 1 else None, + ) + self.half_done.record() + if logits is None: + logits = self.decode._lm_head_shard(self.lm_head, hidden_states) + return logits, hidden_half + + def sample_step( + self, + logits, + hidden_states, + draft_tokens, + step, + positions, + gather_ids, + seq_lens, + kv_lens, + num_accepted, + accepted_tokens, + batch_indices, + *, + hidden_half=None, + ): + """Write the token and lengths once, with optional next-step glue/output rows.""" + rows, draft_len = draft_tokens.shape + last = step + 1 == draft_len + first = step == 0 + if self.use_table and not last: + norm = self.layer.input_layernorm + normed = hidden_states.new_empty((rows, self.hidden_size)) + residual = torch.empty_like(normed) + next_positions = positions.new_empty((rows,)) + torch.cuda.current_stream().wait_event(self.half_done) + glm_draft_ehproj_table( + logits, + self.table_exchange.buffer, + self.table_exchange.flags, + self.decode.mapping.tp_size, + draft_tokens, + step, + hidden_half, + self.token_table, + norm.weight, + norm.variance_epsilon, + normed, + residual, + positions, + next_positions, + gather_ids, + seq_lens, + first, + kv_lens, + 2 if first else 1, + num_accepted if first else None, + draft_len, + None, + [], + None, + None, + None, + None, + None, + ) + return (normed, residual), next_positions, None + glue = None if last else hidden_states.new_empty((rows, self.window[1] - self.window[0])) + next_positions = None if last else positions.new_empty((rows,)) + next_new_tokens = draft_tokens.new_empty((rows, draft_len + 1)) if last else None + glm_draft_sample_glue( + logits, + self.sample_exchange.buffer, + self.sample_exchange.flags, + self.decode.mapping.tp_size, + draft_tokens, + step, + None if last else self.embed_table, + None if last else hidden_states, + None if last else self.enorm_weight, + None if last else self.hnorm_weight, + 0.0 if last else self.layer.enorm.variance_epsilon, + 0 if last else self.window[0], + 0 if last else self.window[1], + glue, + None if last else positions, + next_positions, + None if last else gather_ids, + seq_lens, + first, + kv_lens, + 2 if first else 1, + num_accepted if first else None, + draft_len, + # Current DSA metadata owns these derived outputs and refreshes them below. + None, + [], + None, + None, + None, + None, + None, + accepted_tokens if last else None, + batch_indices if last else None, + num_accepted if last else None, + next_new_tokens, + ) + return glue, next_positions, next_new_tokens + + +@torch.no_grad() +def build_small_batch_draft( + decode: DecodeTarget, layer, lm_head, weights: DraftWeights +) -> SmallBatchDraftKernels | None: + """Construct the helper before KV profiling, with group-uniform setup failure.""" + if decode.mapping.tp_size not in (4, 8): + return None + comm = _get_mnnvl_workspace_comm(decode.mapping) + eligible = False + draft_len = None + try: + spec = decode.core.model_config.spec_config + draft_len = None if spec is None else spec.max_draft_len + eligible = bool(DRAFT_SHAPES.get((decode.mapping.tp_size, draft_len), 0)) and ( + _static_layout_supported(decode, layer, lm_head, weights) + ) + except Exception: + # A rank-local layout/inspection failure must still reach the vote. + pass + if not _all_succeeded(comm, eligible): + return None + # The collective workspace completes on all ranks before local allocations. + table_exchange = ( + DraftExchangeWorkspace.for_ehproj_table(decode.mapping) + if decode.mapping.tp_size == 8 + else None + ) + sample_exchange = DraftExchangeWorkspace.for_sample_glue(decode.mapping) + error, state = None, None + try: + state = SmallBatchDraftKernels( + decode, layer, lm_head, weights, sample_exchange, table_exchange, draft_len=draft_len + ) + torch.cuda.synchronize() + except Exception as exc: + error = exc + if not _all_succeeded(comm, error is None): + raise RuntimeError("small-batch MTP setup failed on at least one rank") from error + return state diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/_speculative.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/_speculative.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/_speculative.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/_speculative.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/_weight_loader.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/_weight_loader.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/_weight_loader.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/_weight_loader.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/modeling.py similarity index 75% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/modeling.py index af2fb54e448b..454927b929b4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/modeling.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GLM-5.2-NVFP4 on 16 Rubin GPUs (sm107), tp16: model construction and forward. +"""GLM-5.2-NVFP4 on Rubin (sm107): model construction and forward. This target serves the generation side of a disaggregated deployment. Its context server hands over its MLA KV cache and DSA indexer K cache. The target @@ -15,7 +15,7 @@ from __future__ import annotations -import contextlib +import os from dataclasses import dataclass from typing import Any, List, Literal, Optional, Tuple @@ -29,7 +29,6 @@ ) from tensorrt_llm._torch.attention.backends.sparse.dsa.module import forward_generation_sparse_attn from tensorrt_llm._torch.attention.backends.sparse.dsa.params import DSABackendForwardArgs -from tensorrt_llm._torch.attention.backends.sparse.params import MTPIndexShareMetadata from tensorrt_llm._torch.distributed import AllReduceFusionOp, AllReduceParams from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl from tensorrt_llm._torch.model_config import ModelConfig @@ -42,31 +41,33 @@ get_data_or_none, to_trtllm_gen_act_type, ) -from tensorrt_llm._torch.speculative import MTPEagleWorker, SpecMetadata -from tensorrt_llm._torch.speculative.eagle3 import _reset_mtp_index_share -from tensorrt_llm._torch.speculative.interface import ( - _FORCE_ACCEPT_RNG_COUNTER_STRIDE, - _FORCE_ACCEPT_RNG_POOL_SIZE, - _FORCE_ACCEPT_RNG_SLOT_STRIDE, -) +from tensorrt_llm._torch.speculative import MTPEagleWorker, SpecMetadata, get_num_extra_kv_tokens from tensorrt_llm._torch.utils import EventType from . import weights as _weights +from ._full_context import FullContextIndexer from ._model import Glm52Model from ._model import setup_aliases as _setup_aliases +from ._mtp_worker import MTPDraftWorker +from ._small_batch_draft import DRAFT_SHAPES, build_small_batch_draft from ._speculative import Glm52SpecDecOneEngineForCausalLM from .ops.attention.mla_kv_half_append import mla_kv_half_append -from .ops.attention.sparse_mla_flash_decode import sparse_mla_flash_decode +from .ops.attention.sparse_mla_flash_decode import ( + sparse_mla_flash_decode, + sparse_mla_flash_decode_latent, +) from .ops.comm._mnnvl_workspace import ( LamportExchangeWorkspace, hop_supported, hop_workspace, publish_target, ) +from .ops.comm.intra_node_reduce import IpcCollectives from .ops.comm.mnnvl_allgather import mnnvl_allgather from .ops.comm.mnnvl_hop import mnnvl_hop from .ops.comm.mnnvl_moe_finalize_hop import mnnvl_moe_finalize_hop from .ops.dsa._fp8_paged_mqa_logits import num_sms +from .ops.dsa._full_context_topk import streaming_plan from .ops.dsa._hist_topk import ( MAX_SHARD, TOPK, @@ -82,10 +83,20 @@ from .ops.dsa.paged_mqa_logits_schedule import paged_mqa_logits_schedule from .ops.dsa.stash_last_accepted_rows import stash_last_accepted_rows from .ops.embedding.fused_embedding import fused_embedding +from .ops.gemm._low_m_splitk import ( + bmm_prewait_tactic, + device_limits, + one_wave_tactic, + publish_supported, + publish_tactic, + run_splitk_bmm, + splitk_default_tactic, +) +from .ops.gemm._low_m_splitk_fused import Q_B_PLANS, q_b_plan from .ops.gemm.dense_mlp_small_m import INTER_PER_RANK, dense_mlp_small_m -from .ops.gemm.draft_hidden_half import draft_hidden_half +from .ops.gemm.draft_hidden_half import draft_hidden_half, hidden_half_tactic from .ops.gemm.fused_a_gemm import INDEXER_OUT_FEATURES as FUSED_A_INDEXER_OUT_FEATURES -from .ops.gemm.fused_a_gemm import MAX_TOKENS as FUSED_A_MAX_TOKENS +from .ops.gemm.fused_a_gemm import NATIVE_MAX_TOKENS as FUSED_A_MAX_TOKENS from .ops.gemm.fused_a_gemm import OUT_FEATURES as FUSED_A_OUT_FEATURES from .ops.gemm.fused_a_gemm import fused_a_gemm from .ops.gemm.indexer_q_tail import indexer_q_tail @@ -107,21 +118,21 @@ from .ops.spec.draft_glue_rmsnorm_cat import glm_draft_glue_rmsnorm_cat from .ops.spec.draft_sample_glue import glm_draft_sample_glue from .ops.spec.gather_last_accepted_rows import glm_gather_last_accepted_rows -from .ops.spec.verify_head_accept import glm_verify_head_accept +from .ops.spec.verify_head_sample import VerifyHeadSampleState # The deployment this target is written for; route() selects it on the same # facts, and the constructor states them again for a direct construction. _COMPUTE_CAPABILITY = (10, 7) -_TP = 16 +_SUPPORTED_TP = (4, 8, 16) # One Lamport buffer of the verify logits' all-gather: every rank's BF16 vocabulary shard of -# up to 8 rows, with room to spare. +# up to 24 rows, with room to spare. _LOGITS_EXCHANGE_BYTES = 16 << 20 class PrefillTarget: """Forwards that carry context rows: the decoder stack the context server runs.""" - def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: + def __init__(self, core: "Glm52Nvfp4Sm107LowLatencyCore") -> None: self.core = core def forward( @@ -191,99 +202,54 @@ class _DraftState: class DecodeTarget: - """Generation-only forwards: the serving step of the generation server. - - Routed to only when the batch has no context rows, which the core decides - per step; nothing in this class reads the phase back. - - A step of up to 8 rows (one request and its 7 draft tokens at c1) runs on - the target's kernels, layer by layer: - - - The embedding is one lookup in the whole BF16 table every rank holds. - - Attention: q_a, the compressed KV and its rope part come from one GEMM over - the stored kv_a_proj_with_mqa weight. q_b, with k_b absorbed and the q_a - norm folded, writes the RoPE'd FP8 query in its epilogue, while on the aux - stream one kernel normalizes, ropes and appends the KV half to the FP8 - cache. The flash decode attends over the 2048 selected rows at 4 heads with - v_b folded in, and o_proj's epilogue publishes this rank's shard into the - all-reduce workspace. - - The DSA indexer (on the layers that run one) shares fused-A: its K and - weight rows are fused-A's last 160 columns. Its Q comes from wq_b (q_a norm - folded) with an epilogue that ropes, quantizes and scales; its K tail - normalizes, ropes, quantizes and writes the indexer K cache on a side - stream. Each rank scores its 1/16 shard of the context's keys, and the ranks - agree on the exact top-2048 by exchanging histograms. The layers without an - indexer reuse the last selection, which the flash decode maps to cache rows - itself and starts gathering before its wait. - - Every all-reduce is a one-shot MNNVL hop that adds the residual and applies - the following RMSNorm: post_attention_layernorm after attention, and after - the MLP or MoE the next layer's input_layernorm (the model's final norm - after the last layer). - - The dense layers 0-2 run this rank's 768 intermediate columns as two - weight-streaming kernels. - - The MoE layers run the shared expert as two split-K GEMMs around SwiGLU on - the aux stream, and the routed experts as one trtllm-gen call that runs the - router GEMV in its routing kernel, joins the NVFP4 input (quantized on its - own stream) right before FC1, runs FC2 on the K = 128 kernel and leaves the - finalize step to the post-MoE hop. - - The verify logits come from the LM head on this rank's vocabulary shard - (split-K GEMM) and one MNNVL all-gather that widens them to FP32. - - A larger batch runs the built-in modules in the same order. - - The MTP draft steps of a one-request step run here too (``draft_first_step``, - ``draft_next_step``, ``draft_head``, ``draft_ehproj_table``, ``draft_sample_glue``), - on state that ``build_draft`` makes once: - - - Draft step 0 runs on the verify rows. One kernel writes this rank's window of - eh_proj's input (embedding gather, enorm, hnorm, concat, column split), eh_proj - runs on it as a split-K GEMM, and its hop applies input_layernorm with a zero - residual. The MTP layer's attention is a verify full-indexer layer's; each - request's top-k of its last accepted row is kept, as KV pool rows of the MTP - layer, for the later steps (on a side stream). After the post-attention hop each - request keeps the row of its last accepted token. - - The later steps take eh_proj's output and its input_layernorm from the previous - step's table kernel. They run no indexer: they attend to draft step 0's - selection, or, with a context short enough that the built-in metadata skips the - indexer, to every position up to their own. - - Every step's MoE runs on that one row: the router GEMV and the low-M BF16 experts, - the shared expert on the aux stream, and the hop that adds both and the residual - and applies shared_head.norm. - - The draft LM head runs on this rank's vocabulary shard. Before every step but the - last, the next step's eh_proj hidden half starts first on the aux stream, and - ``glm_draft_ehproj_table`` exchanges the ranks' argmax and writes the draft - token, the next step's eh_proj output from a token table built at load, its - positions and the sequence and KV lengths. The last step's - ``glm_draft_sample_glue`` writes its token and the next-new-tokens rows. + """Generation kernels selected from the actual row count and TP mapping. + + TP16 uses the one-request histogram indexer, folded MLA output, MNNVL hops, + and inline table-based MTP loop. TP8 uses full-context indexing, folded + 8-head MLA, MNNVL hops and a small-batch table-based MTP loop. + TP4 uses full-context indexing, 16-head + latent MLA followed by batched v_b, IPC collectives, and a small-batch MTP + loop. BF16 dense layers use the fused TP16 kernel or owned split-K GEMMs. + + Both paths fuse the router into the NVFP4 expert call, overlap the shared + expert with routed experts, and gather vocabulary shards from an owned + low-M LM-head projection. Unsupported batches select the built-in layer + stack before writing attention caches. """ - def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: + def __init__(self, core: "Glm52Nvfp4Sm107LowLatencyCore") -> None: self.core = core config = core.model_config.pretrained_config self.mapping = core.model_config.mapping self.hidden_size = config.hidden_size self.index_topk = config.index_topk layers = core.layers[: core.num_hidden_layers] + self._v_b_tactics = {} for layer in layers: assert layer.allreduce is not None and layer.next_layer_layernorm is not None self._check_attention(layer) if isinstance(layer.mlp, GatedMLP): assert ( - layer.mlp_tp_size == _TP + layer.mlp_tp_size == self.mapping.tp_size and not layer.disable_attn_allreduce and not layer.fusion_config.PRE_MLP_FUSION and layer.fusion_config.POST_MLP_FUSION and layer.post_attention_layernorm.nvfp4_scale is None ), f"layer {layer.layer_idx}: unexpected dense-layer configuration" + if self.mapping.tp_size == 8: + assert layer.mlp.gate_up_proj.weight.shape == (3072, self.hidden_size) + assert layer.mlp.down_proj.weight.shape == (self.hidden_size, 1536) else: assert ( layer.disable_attn_allreduce and layer.fusion_config.PRE_MOE_FUSION and layer.fusion_config.POST_MOE_FUSION - and layer.mapping.is_multi_node() ), f"layer {layer.layer_idx}: unexpected MoE-layer configuration" self._check_shared_expert(layer) + if self.mapping.tp_size == 8: + shared = layer.mlp.shared_experts + assert shared.gate_up_proj.weight.shape == (512, self.hidden_size) + assert shared.down_proj.weight.shape == (self.hidden_size, 256) backend = getattr(layer.mlp.experts, "backend", layer.mlp.experts) assert ( backend.has_nvfp4 @@ -303,7 +269,9 @@ def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: # The key-sharded top-k exchanges histograms over the TP group's MNNVL workspace # (created here, on every rank in the same order). self._hist_workspace = ( - get_hist_topk_workspace(self.mapping) if hist_topk_supported(self.mapping) else None + get_hist_topk_workspace(self.mapping) + if self.mapping.tp_size == 16 and hist_topk_supported(self.mapping) + else None ) self._hist_buffers = None self._logits_workspace = LamportExchangeWorkspace(self.mapping, _LOGITS_EXCHANGE_BYTES) @@ -335,6 +303,61 @@ def __init__(self, core: "Glm52Nvfp4Sm107Tp16Core") -> None: # step 0 of the same step runs its indexer on that shard. self._keys_sharded = False self._draft: Optional[_DraftState] = None + self._small_batch_draft = None + self._ipc = None + self._ipc_reason = "" + self._full_context = None + if self.mapping.tp_size == 4: + workspaces = [ + collective.workspace + for layer in core.layers + for collective in (layer.allreduce, layer.moe_allreduce) + if collective is not None + ] + self._ipc, self._ipc_reason = IpcCollectives.bind( + self.mapping, workspaces, max_rows=24, hidden=self.hidden_size + ) + if self.mapping.tp_size in (4, 8): + self._full_context = FullContextIndexer(self, core._decode_context_capacity, device) + if self.mapping.tp_size == 8: + self._prepare_tp8_rows(device) + + def _prepare_tp8_rows(self, device: torch.device) -> None: + """Check the complete TP8 projection/collective route before profiling writes KV.""" + sms, smem = device_limits(device.index) + for rows in (1, 2, 4, 8): + # Dense, shared-expert and first-draft eh_proj projections. + for n, k in ((3072, 6144), (6144, 1536), (512, 6144), (6144, 256)): + splitk_default_tactic(rows, n, k, sm_count=sms, smem_capacity=smem) + head = splitk_default_tactic(rows, 19360, 6144, sm_count=sms, smem_capacity=smem) + one_wave_tactic(head, rows, 19360, 6144, sm_count=sms, smem_capacity=smem) + if rows <= 2: + hidden_half_tactic( + rows, + 768, + self.hidden_size, + in_kernel_norm=rows == 1, + sm_count=sms, + smem_capacity=smem, + ) + projected = publish_tactic( + splitk_default_tactic(rows, 6144, 2048, sm_count=sms, smem_capacity=smem), + smem_capacity=smem, + ) + for plan in Q_B_PLANS: + tactic = q_b_plan(plan, rows, 4608, 2048, device.index) + if plan == "beside_topk": + reserve = ( + ((4608 + tactic.mma_m - 1) // tactic.mma_m) + * ((rows + tactic.mma_n - 1) // tactic.mma_n) + * tactic.split_k + ) + streaming_plan(rows, self._full_context.capacity, sms, reserve) + self._fast_rows[rows] = publish_supported( + rows, self.hidden_size, projected, smem_capacity=smem + ) and hop_supported(rows, self.hidden_size, self.mapping.tp_size) + # Allocate the largest hop before capture; each launch still asks for its current view. + hop_workspace(self.mapping, 8, self.hidden_size) def _check_attention(self, layer) -> None: """Assert the layer's MLA weights and KV cache are the layout the kernels take.""" @@ -358,6 +381,22 @@ def _check_attention(self, layer) -> None: and mla.o_proj.weight.dtype == torch.bfloat16 ), f"layer {layer.layer_idx}: unexpected MLA weights or KV cache layout" assert layer.input_layernorm.nvfp4_scale is None, "kv_a_proj is not NVFP4" + if self.mapping.tp_size == 4: + assert mla.num_heads_tp == 16 and mla.kv_lora_rank == 512 and mla.v_head_dim == 256 + for rows in (1, 2, 3, 6, 12, 18): + tactic = bmm_prewait_tactic(rows, mla.v_head_dim, mla.kv_lora_rank) + assert tactic is not None, f"no v_b tactic for {rows} rows" + self._v_b_tactics[rows] = tactic + elif self.mapping.tp_size == 8: + assert ( + mla.num_heads_tp == 8 + and mla.q_lora_rank == 2048 + and mla.kv_lora_rank == 512 + and mla.qk_rope_head_dim == 64 + and mla.v_head_dim == 256 + and mla.v_b_proj.shape == (8, 256, 512) + and mla.o_proj.weight.shape == (self.hidden_size, 2048) + ), f"layer {layer.layer_idx}: unexpected TP8 MLA geometry" @staticmethod def _check_shared_expert(layer) -> None: @@ -428,7 +467,7 @@ def build_draft(self, layer, lm_head, weights: _weights.DraftWeights) -> None: indexer does under MTP index share; without that share, or where the kernels do not serve one row, nothing is built and the built-in draft loop runs. """ - if self._draft is not None: + if self.mapping.tp_size != 16 or self._draft is not None: return core = self.core config = core.model_config.pretrained_config @@ -531,11 +570,22 @@ def build_draft(self, layer, lm_head, weights: _weights.DraftWeights) -> None: ) def _fast(self, rows: int) -> bool: - """Whether every kernel of the step serves this many rows (decided once per shape).""" + """Decide the complete kernel route before any attention cache writes.""" + if torch.compiler.is_compiling(): + return False + if self.mapping.tp_size == 8: + return self._fast_rows.get(rows, False) + if self.mapping.tp_size == 4: + return ( + rows in self._v_b_tactics + and self._ipc is not None + and self._ipc.supported(rows) + and self._ipc.supported(rows, finalize=True) + ) fast = self._fast_rows.get(rows) if fast is None: fast = ( - rows <= FUSED_A_MAX_TOKENS + 1 <= rows <= 8 and rows <= FUSED_ROUTER_MAX_TOKENS and bool( torch.ops.trtllm.glm_dense_mlp_small_m_supported( @@ -547,6 +597,21 @@ def _fast(self, rows: int) -> bool: self._fast_rows[rows] = fast return fast + def supports_small_batch_draft(self, metadata, first_rows, batch_size, layer) -> bool: + tp = self.mapping.tp_size + return ( + tp in (4, 8) + and batch_size > 0 + and first_rows % batch_size == 0 + and batch_size <= DRAFT_SHAPES.get((tp, first_rows // batch_size - 1), 0) + and self._full_context is not None + and self._full_context.supports(metadata, first_rows) + and self._fast(first_rows) + and self._fast(batch_size) + and layer.layer_idx in self._q_b_absorbed + and layer.layer_idx in self._fused_a_indexer_weight + ) + def forward( self, attn_metadata: AttentionMetadata, @@ -562,7 +627,10 @@ def forward( inputs_embeds = fused_embedding(input_ids, core.embed_tokens.weight) residual = inputs_embeds hidden_states = layers[0].input_layernorm(inputs_embeds) - if not self._fast(hidden_states.shape[0]): + if not self._fast(hidden_states.shape[0]) or ( + self._full_context is not None + and not self._full_context.supports(attn_metadata, hidden_states.shape[0]) + ): self._keys_sharded = False for layer in layers: hidden_states, residual = self._builtin_layer( @@ -573,16 +641,37 @@ def forward( self._keys_sharded = sharded for layer in layers: attn_out = self._attention(layer, position_ids, hidden_states, attn_metadata, sharded) - hidden_states, residual = self._post_attention_hop(layer, attn_out, residual) + quantized = None + if self._ipc is not None and not isinstance(layer.mlp, GatedMLP): + hidden_states, residual, packed, scales = self._post_attention_hop( + layer, attn_out, residual, premoe=True + ) + quantized = (packed, scales) + else: + hidden_states, residual = self._post_attention_hop(layer, attn_out, residual) if isinstance(layer.mlp, GatedMLP): mlp = layer.mlp - # This rank's partial sum; the hop adds the 16 ranks'. - partial = dense_mlp_small_m( - hidden_states, mlp.gate_up_proj.weight, mlp.down_proj.weight - ) + # This rank's partial sum; the hop adds the TP ranks'. + if self.mapping.tp_size == 16: + partial = dense_mlp_small_m( + hidden_states, mlp.gate_up_proj.weight, mlp.down_proj.weight + ) + else: + gate_up = low_m_gemm( + hidden_states, + mlp.gate_up_proj.weight, + hidden_states.new_empty( + hidden_states.shape[0], mlp.gate_up_proj.weight.shape[0] + ), + ) + partial = low_m_gemm( + swiglu(gate_up), + mlp.down_proj.weight, + hidden_states.new_empty(hidden_states.shape[0], self.hidden_size), + ) hidden_states, residual = self._hop(layer.next_layer_layernorm, partial, residual) else: - hidden_states, residual = self._moe(layer, hidden_states, residual) + hidden_states, residual = self._moe(layer, hidden_states, residual, quantized) return hidden_states def logits(self, lm_head, hidden_states: torch.Tensor) -> Optional[torch.Tensor]: @@ -590,11 +679,13 @@ def logits(self, lm_head, hidden_states: torch.Tensor) -> Optional[torch.Tensor] do not serve: this rank's vocabulary shard from the LM head, then the all-gather.""" rows = hidden_states.shape[0] weight = lm_head.weight - if not self._fast(rows) or weight.shape[0] * self.mapping.tp_size != lm_head.num_embeddings: + if ( + not 1 <= rows <= 24 + or torch.compiler.is_compiling() + or weight.shape[0] * self.mapping.tp_size != lm_head.num_embeddings + ): return None - shard = low_m_gemm( - hidden_states.contiguous(), weight, hidden_states.new_empty(rows, weight.shape[0]) - ) + shard = self._lm_head_shard(lm_head, hidden_states) return mnnvl_allgather( shard, self._logits_workspace.comm_buffer, @@ -602,6 +693,15 @@ def logits(self, lm_head, hidden_states: torch.Tensor) -> Optional[torch.Tensor] self.mapping.tp_size, ) + def _lm_head_shard(self, lm_head, hidden_states): + weight = lm_head.weight + return low_m_gemm( + hidden_states.contiguous(), + weight, + hidden_states.new_empty(hidden_states.shape[0], weight.shape[0]), + one_wave=True, + ) + def drafts(self, rows: int) -> bool: """Whether this step's MTP draft steps run here: the draft state exists, this step's verify forward put the indexer on this rank's key shard (one request; draft step 0 @@ -955,7 +1055,7 @@ def _shard_indexer_keys(self, attn_metadata) -> bool: or buffers["tokens_per_block"] != tokens_per_block or buffers["logits"].device != device ): - rows = FUSED_A_MAX_TOKENS + rows = 8 sms = num_sms(device.index) self._hist_buffers = dict( capacity=capacity, @@ -997,6 +1097,8 @@ def _attention( sharded, pool_rows: Optional[torch.Tensor] = None, positions: Optional[torch.Tensor] = None, + *, + draft_state=None, ): """MLA with the DSA indexer over the generation rows: the attention output ahead of o_proj, [rows, heads * v_head_dim]. @@ -1017,11 +1119,24 @@ def _attention( local_layer = mqa.get_local_layer_idx(attn_metadata) # Asked per step: the executor re-creates the cache manager after its memory profiling. kv_pool = manager.get_unique_primary_pool() + full_context = self._full_context + if full_context is not None and full_context.reuse_draft_selection( + indexer, attn_metadata, draft_state + ): + positions = draft_state.topk_stash[: attn_metadata.num_seqs] reuses_selection = pool_rows is not None or positions is not None - own_indexer = sharded and indexer is not None and not reuses_selection + own_indexer = ( + (sharded or full_context is not None) and indexer is not None and not reuses_selection + ) # A layer that attends to the selection of the indexer layer before it gathers its old KV # rows before the flash decode's wait. - early_gather = sharded and indexer is None and not reuses_selection + early_gather = ( + (sharded or (full_context is not None and attn_metadata.num_seqs <= 2)) + and indexer is None + and not reuses_selection + ) + if full_context is not None and indexer is None: + positions = attn_metadata.shared_topk_indices[:num_tokens] weight = ( self._fused_a_indexer_weight[layer.layer_idx] @@ -1029,11 +1144,11 @@ def _attention( else mla.kv_a_proj_with_mqa.weight ) if reuses_selection: - l2_prefetch = self._draft.l2_after_fused_a + l2_prefetch = (self._draft if draft_state is None else draft_state).l2_after_fused_a elif early_gather: # Also the selection the flash decode reads first. l2_prefetch = self._l2_after_fused_a[layer.layer_idx] + [ - self._hist_buffers["topk"][:num_tokens] + self._hist_buffers["topk"][:num_tokens] if sharded else positions ] else: l2_prefetch = self._l2_after_fused_a[layer.layer_idx] @@ -1085,23 +1200,46 @@ def kv_half(trigger_at_entry=False): (hidden_states.shape[0], mla.num_heads_tp * mla.v_head_dim) ) q_heads = quant_q.view(num_tokens, mla.num_heads_tp, head_dim) - if sharded: + if sharded or full_context is not None: if own_indexer: def select(): - self._select_rows( - indexer, kv_a, latent_end, q_a, position_ids, attn_metadata, num_tokens, mla - ) + nonlocal positions + if full_context is not None: + positions = full_context.select( + layer, + kv_a, + q_a, + position_ids, + attn_metadata, + draft_state, + beside_topk=kv_half_then_q_b, + ) + else: + self._select_rows( + indexer, + kv_a, + latent_end, + q_a, + position_ids, + attn_metadata, + num_tokens, + mla, + ) def kv_half_then_q_b(): kv_half() - # Queued after the top-k: 18 CTAs leave the SMs to the indexer beside it. + # The q_b launch shares SMs with the reserved TopK grid. q_b("beside_topk") - # The indexer on this stream; the KV half and q_b beside it on the aux stream. - maybe_execute_in_parallel( - select, kv_half_then_q_b, mla.ln_events[0], mla.ln_events[1], mla.aux_stream - ) + if full_context is not None: + # The full-context path queues q_b only once logits are ready, + # leaving the earlier full-query scoring pass all of the SMs. + select() + else: + maybe_execute_in_parallel( + select, kv_half_then_q_b, mla.ln_events[0], mla.ln_events[1], mla.aux_stream + ) else: # With the early gather, q_b at split-K 4 leaves room for the flash decode, which # the KV half lets launch at its entry. @@ -1112,6 +1250,38 @@ def kv_half_then_q_b(): mla.ln_events[1], mla.aux_stream, ) + if self.mapping.tp_size == 4 and full_context is not None: + page_index_scale, layer_offset = manager.get_primary_pool_page_index_params( + local_layer + ) + latent = hidden_states.new_empty((num_tokens, mla.num_heads_tp, mla.kv_lora_rank)) + sparse_mla_flash_decode_latent( + q_heads, + kv_pool, + positions, + kv_lens, + self._bmm1_scale, + self._bmm2_scale, + latent, + seq_len_q, + self.index_topk, + attn_metadata.block_table[: attn_metadata.num_seqs], + manager.tokens_per_block, + page_index_scale, + layer_offset, + early_gather, + True, + ) + tactic = self._v_b_tactics[num_tokens] + run_splitk_bmm( + latent.transpose(0, 1), + mla.v_b_proj, + output.view(num_tokens, mla.num_heads_tp, mla.v_head_dim).transpose(0, 1), + get_env_enable_pdl(), + tactic, + reserve_full_smem=False, + ) + return output if pool_rows is not None: sparse_mla_flash_decode( q_heads, @@ -1268,10 +1438,23 @@ def _select_rows(self, indexer, kv_a, latent_end, q_a, position_ids, attn_metada self.mapping.tp_size, ) - def _post_attention_hop(self, layer, attn_out: torch.Tensor, residual: torch.Tensor): + def _post_attention_hop( + self, layer, attn_out: torch.Tensor, residual: torch.Tensor, *, premoe=False + ): """o_proj published into the all-reduce workspace, then the hop with the residual add and post_attention_layernorm: (normed, residual).""" rows = attn_out.shape[0] + if self._ipc is not None: + shard = low_m_gemm( + attn_out, + layer.self_attn.o_proj.weight, + attn_out.new_empty((rows, self.hidden_size)), + ) + norm = layer.post_attention_layernorm + if premoe: + backend = getattr(layer.mlp.experts, "backend", layer.mlp.experts) + return self._ipc.premoe(shard, residual, norm, backend.fc31_input_scale) + return self._ipc.plain_reduce(norm, shard, residual) lamport, flags = publish_target(self.mapping, rows, self.hidden_size) comm_buffer, buffer_flags = hop_workspace(self.mapping, rows, self.hidden_size) shard = attn_out.new_empty((rows, self.hidden_size)) @@ -1291,24 +1474,35 @@ def _post_attention_hop(self, layer, attn_out: torch.Tensor, residual: torch.Ten def _hop(self, norm, partial: torch.Tensor, residual: torch.Tensor): """The hop of a rank-partial sum with the residual add and `norm`: (normed, residual).""" + if self._ipc is not None: + return self._ipc.plain_reduce(norm, partial, residual) comm_buffer, buffer_flags = hop_workspace(self.mapping, partial.shape[0], self.hidden_size) return mnnvl_hop( partial, residual, norm.weight, norm.variance_epsilon, comm_buffer, buffer_flags, False ) - def _moe(self, layer, hidden_states: torch.Tensor, residual: torch.Tensor): + def _moe(self, layer, hidden_states: torch.Tensor, residual: torch.Tensor, quantized=None): """The MoE block and the post-MoE hop that finalizes it: (normed, residual).""" mlp = layer.mlp backend = getattr(mlp.experts, "backend", mlp.experts) rows = hidden_states.shape[0] - self._quant_fork.record() - with torch.cuda.stream(self._quant_stream): - self._quant_stream.wait_event(self._quant_fork) - x_fp4, x_sf = fp4_quantize( - hidden_states, backend.fc31_input_scale, backend.scaling_vector_size, False, False - ) - self._quant_ready.record() + if quantized is None: + self._quant_fork.record() + with torch.cuda.stream(self._quant_stream): + self._quant_stream.wait_event(self._quant_fork) + x_fp4, x_sf = fp4_quantize( + hidden_states, + backend.fc31_input_scale, + backend.scaling_vector_size, + False, + False, + ) + self._quant_ready.record() + quant_ready = self._quant_ready.cuda_event + else: + x_fp4, x_sf = quantized + quant_ready = 0 routing = extract_routing_params(backend.routing_method) def routed(): @@ -1346,7 +1540,7 @@ def routed(): None, backend.max_num_tokens, False, - self._quant_ready.cuda_event, + quant_ready, hidden_states, mlp.gate.weight, self._router_workspace[layer.layer_idx], @@ -1361,6 +1555,10 @@ def routed(): mlp.aux_stream, ) norm = layer.next_layer_layernorm + if self._ipc is not None: + return self._ipc.finalize( + fc2_output, expert_weights, permuted_rows, shared_output, residual, norm + ) comm_buffer, buffer_flags = hop_workspace(self.mapping, rows, self.hidden_size) return mnnvl_moe_finalize_hop( fc2_output, @@ -1474,13 +1672,14 @@ def _attention_builtin(layer, position_ids, hidden_states, attn_metadata, all_re return mla.o_proj(output, all_reduce_params=all_reduce_params, layer_idx=mla.layer_idx) -class Glm52Nvfp4Sm107Tp16Core(Glm52Model): +class Glm52Nvfp4Sm107LowLatencyCore(Glm52Model): """The decoder stack, with the per-step choice between the two targets.""" def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: super().__init__(model_config) self._prefill: Optional[PrefillTarget] = None self._decode: Optional[DecodeTarget] = None + self._decode_context_capacity = None @torch.no_grad() def post_load_weights(self) -> None: @@ -1505,7 +1704,14 @@ def build_draft(self, layer, lm_head, weights: _weights.DraftWeights) -> None: """Build the decode target's MTP draft-step state (``DecodeTarget.build_draft``).""" if self._decode is None: self.post_load_weights() - self._decode.build_draft(layer, lm_head, weights) + if self.model_config.mapping.tp_size in (4, 8): + decode = self._decode + decode._check_attention(layer) + decode._check_shared_expert(layer) + decode._fold_attention_weights(layer) + decode._small_batch_draft = build_small_batch_draft(decode, layer, lm_head, weights) + else: + self._decode.build_draft(layer, lm_head, weights) def forward( self, @@ -1526,342 +1732,10 @@ def forward( ) -class MTPDraftWorker(MTPEagleWorker): - """The one-model MTP-Eagle worker: the draft loop of a one-request generation step runs - on the decode target's kernels (``DecodeTarget``'s draft steps); every other step runs the - built-in loop. - - The loop keeps the built-in loop's attention-metadata handling: the MTP index-share flags, - the switch to one row per request after draft step 0, and the cleanup. The draft-step - kernels advance the sequence and KV lengths in place, and the draft tokens and the - next-new-tokens rows come from them. - """ - - def __init__( - self, - spec_config, - model_config: ModelConfig, - use_separate_draft_kv_cache: bool, - *, - mapping, - decode_target, - ) -> None: - super().__init__(spec_config, model_config, use_separate_draft_kv_cache, mapping=mapping) - # The core's per-step lookup of the decode target (a bound method rather than the - # core itself, which belongs to the shell and is no submodule of this worker). - self._decode_target = decode_target - # The draft tokens of the last draft loop on the decode target and the next-new-tokens - # rows its last step wrote, for _prepare_next_new_tokens. - self._drafted_rows: Optional[Tuple[torch.Tensor, torch.Tensor]] = None - - def prepare_1st_drafter_inputs( - self, - input_ids, - position_ids, - hidden_states, - accepted_tokens, - attn_metadata, - spec_metadata, - draft_model, - ): - """Without context requests draft step 0's input ids are the accepted-token rows as - they are, flattened, rather than a copy behind an empty context part.""" - if not self.is_mtp_eagle or attn_metadata.num_contexts > 0: - return super().prepare_1st_drafter_inputs( - input_ids, - position_ids, - hidden_states, - accepted_tokens, - attn_metadata, - spec_metadata, - draft_model, - ) - return { - "input_ids": accepted_tokens[:, : spec_metadata.runtime_draft_len + 1].flatten(), - "position_ids": position_ids, - "hidden_states": hidden_states, - "attn_metadata": attn_metadata, - "spec_metadata": spec_metadata, - } - - def _prepare_next_new_tokens( - self, - accepted_tokens, - next_draft_tokens, - batch_indices_cuda, - batch_size, - num_accepted_tokens, - ): - """The rows the last draft-step kernel wrote when these draft tokens are that loop's.""" - drafted, self._drafted_rows = self._drafted_rows, None - if drafted is not None and drafted[0] is next_draft_tokens: - return drafted[1] - return super()._prepare_next_new_tokens( - accepted_tokens, - next_draft_tokens, - batch_indices_cuda, - batch_size, - num_accepted_tokens, - ) - - def _sample_and_accept_draft_tokens_base( - self, - logits, - draft_tokens, - num_contexts, - batch_size, - spec_metadata, - ): - """Strict acceptance and the forced accepted-token count in one kernel - (``glm_verify_head_accept``) for a step of generation requests only, with the built-in - target-token sampling; ragged verification and context requests take the built-in - chain. The forced count follows ``_apply_force_accepted_tokens``: none during eager - CUDA-graph warm-up, the integer part plus one, and with a fractional part the same pool - draw and counter step.""" - if spec_metadata is None or spec_metadata.is_ragged_verify or num_contexts != 0: - return super()._sample_and_accept_draft_tokens_base( - logits, draft_tokens, num_contexts, batch_size, spec_metadata - ) - if logits.dim() == 1: - logits = logits.unsqueeze(0) - tokens_per_request = draft_tokens.shape[-1] + 1 - target_tokens = self._sample_tokens_for_batch(logits, spec_metadata, 0, batch_size) - mode, base, frac = 0, 0, 0.0 - warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() - if self.force_num_accepted_tokens != 0.0 and not warmup: - int_part = int(self.force_num_accepted_tokens) - frac = self.force_num_accepted_tokens - int_part - base = min(int_part + 1, tokens_per_request) - mode = 2 if frac > 0.0 and base < tokens_per_request else 1 - pool = counter = None - if mode == 2: - self._ensure_force_accept_rng_state(logits.device) - pool, counter = self._force_accept_rng_pool, self._force_accept_rng_counter - accepted_tokens = torch.empty( - (batch_size, tokens_per_request), dtype=torch.int, device=logits.device - ) - num_accepted_tokens = glm_verify_head_accept( - target_tokens, - draft_tokens, - 0, - accepted_tokens, - pool, - counter, - mode, - base, - tokens_per_request, - frac, - _FORCE_ACCEPT_RNG_COUNTER_STRIDE, - _FORCE_ACCEPT_RNG_SLOT_STRIDE, - _FORCE_ACCEPT_RNG_POOL_SIZE, - ) - return accepted_tokens, num_accepted_tokens - - def _forward_linear_draft_loop( - self, - inputs, - attn_metadata, - spec_metadata, - draft_model, - draft_kv_cache_manager, - num_contexts, - batch_size, - num_accepted_tokens, - original_all_rank_num_tokens, - accepted_tokens=None, - pre_draft_seq_lens=None, - ): - self._drafted_rows = None - decode = self._drafting_target( - inputs, - attn_metadata, - spec_metadata, - draft_kv_cache_manager, - num_contexts, - batch_size, - num_accepted_tokens, - ) - if decode is None: - return super()._forward_linear_draft_loop( - inputs, - attn_metadata, - spec_metadata, - draft_model, - draft_kv_cache_manager, - num_contexts, - batch_size, - num_accepted_tokens, - original_all_rank_num_tokens, - accepted_tokens, - pre_draft_seq_lens, - ) - draft_tokens = self._draft_on_target( - decode, - inputs, - attn_metadata, - spec_metadata, - draft_kv_cache_manager, - batch_size, - num_accepted_tokens, - ) - return draft_tokens, None - - def _drafting_target( - self, - inputs, - attn_metadata, - spec_metadata, - draft_kv_cache_manager, - num_contexts, - batch_size, - num_accepted_tokens, - ) -> Optional[DecodeTarget]: - """The decode target when its kernels draft this step, else None. - - They draft a generation step of one request whose draft layer shares the target's KV - pool, through the configured draft length with no draft-length schedule (the draft - steps after the first skip the MTP layer's indexer K write: the next step's draft - step 0 writes those positions before it reads them), greedily (argmax drafts), with - nothing that rewrites the drafts (guided decoding, SA, a draft vocabulary map), and - when the decode target's kernels serve the rows. - """ - draft_len = spec_metadata.runtime_draft_len - rows = batch_size * (draft_len + 1) - input_ids = inputs["input_ids"] - position_ids = inputs["position_ids"] - if not ( - self.is_mtp_eagle - and num_contexts == 0 - and batch_size == 1 - and draft_kv_cache_manager is None - and draft_len == self.max_draft_len - and getattr(self.spec_config, "draft_len_schedule", None) is None - and self.guided_decoder is None - and self.sa_enhancer is None - and self._d2t is None - and not spec_metadata.wants_advanced_draft_sampling - and not spec_metadata.is_ragged_verify - and attn_metadata.num_tokens == rows - and inputs["hidden_states"].shape[0] == rows - and input_ids.dtype == torch.int32 - and tuple(input_ids.shape) == (rows,) - and input_ids.is_contiguous() - and position_ids.dtype in (torch.int32, torch.int64) - and tuple(position_ids.shape) == (rows,) - and position_ids.is_contiguous() - and num_accepted_tokens.dtype == torch.int32 - and num_accepted_tokens.is_contiguous() - and spec_metadata.batch_indices_cuda.dtype == torch.int32 - ): - return None - decode = self._decode_target(attn_metadata) - if decode is None or not decode.drafts(rows): - return None - return decode - - def _draft_on_target( - self, - decode: DecodeTarget, - inputs, - attn_metadata, - spec_metadata, - draft_kv_cache_manager, - batch_size, - num_accepted_tokens, - ) -> torch.Tensor: - """The draft loop on the decode target's kernels: the draft tokens [batch, draft_len]. - - Each step's last kernel writes the next step's sequence and KV lengths in place. The - steps after the first run no indexer and read nothing else of the DSA decode metadata, - so the built-in loop's per-step rebuild of it (``update_for_spec_dec``) is left out; - with a context short enough that the metadata skips the indexer, only the next step's - dense selection is refreshed from the new lengths. The host side of draft step 0's - switch to one row per request comes before that. - """ - draft_len = spec_metadata.runtime_draft_len - # The accepted tokens' rows: draft step 0's input ids (prepare_1st_drafter_inputs). - accepted_tokens = inputs["input_ids"].view(batch_size, draft_len + 1) - batch_indices = spec_metadata.batch_indices_cuda[:batch_size] - num_accepted = num_accepted_tokens[:batch_size] - draft_tokens = accepted_tokens.new_empty((batch_size, draft_len)) - positions = inputs["position_ids"] - eh_proj_output = gather_ids = pool_rows = next_new_tokens = None - uses_mtp_index_share = isinstance(attn_metadata, MTPIndexShareMetadata) - with contextlib.ExitStack() as draft_scope: - if uses_mtp_index_share: - attn_metadata.set_in_mtp_draft_loop(True) - attn_metadata.set_mtp_num_accepted(num_accepted_tokens) - draft_scope.callback(_reset_mtp_index_share, attn_metadata) - attn_metadata = draft_scope.enter_context( - self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager) - ) - # The draft loop's copy of the sequence lengths (prepare_for_spec_dec) and the KV - # lengths, which the draft-step kernels advance. - seq_lens = attn_metadata.seq_lens_cuda[:batch_size] - kv_lens = attn_metadata.kv_lens_cuda[:batch_size] - for i in range(draft_len): - if uses_mtp_index_share: - attn_metadata.set_skip_topk(i > 0) - if i == 0: - hidden_states, gather_ids, pool_rows = decode.draft_first_step( - attn_metadata, - inputs["input_ids"], - positions, - inputs["hidden_states"], - batch_indices, - num_accepted, - draft_len, - ) - else: - hidden_states = decode.draft_next_step( - attn_metadata, eh_proj_output, pool_rows, join_selection=i == 1 - ) - last = i == draft_len - 1 - logits, hidden_half = decode.draft_head(hidden_states, next_step=not last) - rewind_accepted = num_accepted if i == 0 else None - if i == 0: - # One row per request from here on, on the host (no context requests to - # turn into generation ones); the kernel below fills seq_lens. - attn_metadata._seq_lens[:batch_size].fill_(1) - attn_metadata.on_update() - if last: - next_new_tokens = decode.draft_sample_glue( - logits, - draft_tokens, - i, - seq_lens, - kv_lens, - rewind_accepted, - accepted_tokens, - batch_indices, - num_accepted, - ) - else: - eh_proj_output, positions = decode.draft_ehproj_table( - logits, - hidden_half, - draft_tokens, - i, - positions, - gather_ids if i == 0 else None, - seq_lens, - kv_lens, - rewind_accepted, - ) - if not last and attn_metadata.skip_indexer_for_gen_reqs: - attn_metadata.prepare_dense_topk_indices( - attn_metadata.kv_lens_cuda, device=True - ) - if i == 0: - attn_metadata.use_spec_decoding = False - self._drafted_rows = (draft_tokens, next_new_tokens) - return draft_tokens - - -class ModelingV2Glm52Nvfp4Sm107Tp16( - Glm52SpecDecOneEngineForCausalLM[Glm52Nvfp4Sm107Tp16Core, PretrainedConfig] +class ModelingV2Glm52Nvfp4Sm107LowLatency( + Glm52SpecDecOneEngineForCausalLM[Glm52Nvfp4Sm107LowLatencyCore, PretrainedConfig] ): - """GLM-5.2-NVFP4, sm107, tp16.""" + """GLM-5.2-NVFP4 on Rubin, with tensor-parallel attention and experts.""" @classmethod def get_preferred_kv_cache_manager_version(cls, pretrained_config: Any = None) -> Literal["V2"]: @@ -1881,14 +1755,20 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: f"this target runs on sm107, not {torch.cuda.get_device_capability()}" ) assert ( - mapping.world_size == _TP - and mapping.tp_size == _TP + mapping.tp_size in _SUPPORTED_TP + and mapping.world_size == mapping.tp_size + and mapping.moe_tp_size == mapping.tp_size + and mapping.cp_size == 1 and mapping.moe_ep_size == 1 and not mapping.enable_attention_dp and mapping.pp_size == 1 - ), f"this target is tp16 without attention DP or expert parallelism, not {mapping}" + ), ( + f"this target uses TP4/TP8/TP16 without attention DP or expert parallelism, not {mapping}" + ) model_config = _weights.normalize_model_config(model_config) - super().__init__(model=Glm52Nvfp4Sm107Tp16Core(model_config), model_config=model_config) + super().__init__( + model=Glm52Nvfp4Sm107LowLatencyCore(model_config), model_config=model_config + ) if ( model_config.spec_config is not None and model_config.spec_config.spec_dec_mode.is_mtp_one_model() @@ -1977,6 +1857,26 @@ def setup_aliases(self) -> None: def post_load_weights(self): super().post_load_weights() + inferred = self.infer_max_seq_len() + requested = self.model.model_config.max_seq_len + allow_long = os.getenv("TLLM_ALLOW_LONG_MAX_MODEL_LEN", "0").lower() in ( + "1", + "true", + "yes", + "y", + ) + net = ( + inferred if requested is None else requested if allow_long else min(requested, inferred) + ) + spec = self.model.model_config.spec_config + extra = ( + 0 if spec is None else 2 * spec.max_total_draft_tokens + get_num_extra_kv_tokens(spec) + ) + self.model._decode_context_capacity = net + extra + if self.model.model_config.mapping.tp_size in (4, 8) and isinstance( + self.spec_worker, MTPDraftWorker + ): + self.spec_worker._verify_sampler = VerifyHeadSampleState(self.lm_head.weight.device) self.model.post_load_weights() draft_weights, self._draft_weights = self._draft_weights, None if draft_weights is not None: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/_cpp_fakes.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/_cpp_fakes.py similarity index 64% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/_cpp_fakes.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/_cpp_fakes.py index c8f37d63dc7e..310bc5d415b2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/_cpp_fakes.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/_cpp_fakes.py @@ -150,3 +150,101 @@ def _glm_mnnvl_moe_finalize_hop( @torch.library.register_fake("trtllm::glm_mnnvl_allgather") def _glm_mnnvl_allgather(input, comm_buffer, buffer_flags, group_size): return input.new_empty((input.shape[0], input.shape[1] * group_size), dtype=torch.float32) + + +@torch.library.register_fake("trtllm::glm_sparse_mla_flash_decode") +def _glm_sparse_mla_flash_decode( + q, + kv_pool, + indices, + seq_lens_kv, + mla_bmm1_scale, + mla_bmm2_scale, + v_b_proj, + output, + seq_len_q, + top_k, + block_table=None, + tokens_per_block=0, + page_index_scale=0, + layer_offset=0, + gather_before_wait=False, +): + return None + + +@torch.library.register_fake("trtllm::glm_sparse_mla_flash_decode_latent") +def _glm_sparse_mla_flash_decode_latent( + q, + kv_pool, + indices, + seq_lens_kv, + mla_bmm1_scale, + mla_bmm2_scale, + output, + seq_len_q, + top_k, + block_table=None, + tokens_per_block=0, + page_index_scale=0, + layer_offset=0, + gather_before_wait=False, + trigger_after_wait=False, +): + return None + + +@torch.library.register_fake("trtllm::glm_intra_node_reduce") +def _glm_intra_node_reduce( + input, + residual, + gamma, + epsilon, + table, + flags, + layout, + rank, + nranks, + buffer_bytes, + trigger_completion_at_end, +): + return torch.empty_like(residual), torch.empty_like(residual) + + +@torch.library.register_fake("trtllm::glm_intra_node_moe_finalize") +def _glm_intra_node_moe_finalize( + fc2, + weights, + indices, + shared, + residual, + gamma, + epsilon, + table, + flags, + layout, + rank, + nranks, + buffer_bytes, +): + return torch.empty_like(residual), torch.empty_like(residual) + + +@torch.library.register_fake("trtllm::glm_pre_moe_fp4_linear") +def _glm_pre_moe_fp4_linear( + input, residual, gamma, epsilon, scale, table, flags, layout, rank, nranks, buffer_bytes +): + rows, hidden = residual.shape + return ( + torch.empty_like(residual), + torch.empty_like(residual), + residual.new_empty((rows, hidden // 2), dtype=torch.uint8), + residual.new_empty((rows, hidden // 16), dtype=torch.uint8), + ) + + +@torch.library.register_fake("trtllm::glm_verify_head_sample_tokens") +def _glm_verify_head_sample_tokens( + logits, temperatures, top_ks, top_ps, min_ps, seed, offset, row_counters, workspace +): + return logits.new_empty((logits.shape[0],), dtype=torch.int32) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/attention/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/attention/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/mla_kv_half_append.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/attention/mla_kv_half_append.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/mla_kv_half_append.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/attention/mla_kv_half_append.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/sparse_mla_flash_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/attention/sparse_mla_flash_decode.py similarity index 64% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/sparse_mla_flash_decode.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/attention/sparse_mla_flash_decode.py index fea70f687292..218200e79af8 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/attention/sparse_mla_flash_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/attention/sparse_mla_flash_decode.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""DSA sparse MLA generation attention at 4 query heads per rank, with the v_b BMM folded in.""" +"""DSA sparse MLA generation attention at 4 or 8 query heads per rank, with the v_b BMM folded in.""" from __future__ import annotations @@ -19,6 +19,9 @@ #: Query heads per rank: GLM-5.2's 64 MLA heads over 16 ranks. NUM_HEADS = 4 +TP8_NUM_HEADS = 8 +LATENT_NUM_HEADS = 16 +LATENT_MAX_TOKENS = 22 #: The latent row: kv_lora_rank 512 followed by qk_rope_head_dim 64. HEAD_DIM_QK = 576 KV_LORA_RANK = 512 @@ -120,11 +123,34 @@ class GlmSparseMlaFlashDecode(OpWrapper): gather_before_wait=True, ), ), + *( + Cell( + why=f"TP8 folded flash: {seq_len_q} query rows x {len(kv_lens)} requests, early={early}", + spec=_cell( + heads=TP8_NUM_HEADS, + seq_len_q=seq_len_q, + kv_lens=kv_lens, + remap=remap, + layer_offset=7, + gather_before_wait=early, + holes=True, + ), + ) + for seq_len_q, kv_lens, remap, early in ( + (1, (1500,), False, False), + (1, (9000, 17003), True, True), + (4, (1500,), True, False), + (4, (9000, 17003), True, True), + (8, (1500,), True, False), + (8, (20011,), True, False), + (8, (20011,), True, True), + ) + ), ) note = """ Metadata consumed: none -- every input is an argument. What the target passes: - - `q`: the [T, 4, 576] E4M3 query that q_b's fused epilogue writes (absorbed q_nope, + - `q`: the [T, 4 or 8, 576] E4M3 query that q_b's fused epilogue writes (absorbed q_nope, roped q_pe); `seq_len_q` the query tokens per sequence (8 on verify and MTP draft step 0, 1 on draft steps 1-6). - `kv_pool`: the DSA cache manager's `get_unique_primary_pool()`, the base that @@ -255,7 +281,7 @@ def reference( with true_fp32_matmul(): scores = torch.einsum("thd,tjd->thj", queries, keys) * mla_bmm1_scale[1].float() scores = scores.masked_fill(~valid[:, None, :], float("-inf")) - scores = scores.view(num_tokens, NUM_HEADS, _SPLITS, _SLOTS_PER_SPLIT) + scores = scores.view(num_tokens, q.shape[1], _SPLITS, _SLOTS_PER_SPLIT) split_max = scores.amax(dim=-1) empty = split_max == float("-inf") p = torch.exp2(scores - torch.where(empty, 0.0, split_max)[..., None]) @@ -278,7 +304,7 @@ def reference( attn = attn.to(torch.bfloat16).float() with true_fp32_matmul(): out = torch.einsum("thd,hvd->thv", attn, v_b_proj.float()) - return out.reshape(num_tokens, NUM_HEADS * V_HEAD_DIM).to(torch.bfloat16) + return out.reshape(num_tokens, q.shape[1] * V_HEAD_DIM).to(torch.bfloat16) def is_valid( self, @@ -361,3 +387,202 @@ def is_valid( sparse_mla_flash_decode = GlmSparseMlaFlashDecode() + + +class GlmSparseMlaFlashDecodeLatent(OpWrapper): + """TP4 sparse MLA: sixteen query heads, explicit BF16 latent output [T,16,512]. + + For up to eight query tokens the top-k is split over sixteen CTAs; for nine to + twenty-two, over eight CTAs. The v_b projection is a separate consumer. With + trigger_after_wait=True that consumer must stage only its weights before its + own dependency wait and read this output after the wait. The default is False. + Index, cache and gather-before-wait semantics match the folded four-head op. + """ + + ARCHS = frozenset({Arch.SM_107}) + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"TP4 draft batch {batch}: sixteen latent heads, global rows", + spec=_cell(seq_len_q=1, kv_lens=(4099,) * batch, remap=False, layer_offset=78), + ) + for batch in (1, 2, 3) + ), + *( + Cell( + why=f"TP4 verify batch {batch}: {6 * batch} queries, remapped pool rows", + spec=_cell( + seq_len_q=6, + kv_lens=(4101,) * batch, + remap=True, + layer_offset=6, + gather_before_wait=batch == 2, + ), + ) + for batch in (1, 2, 3) + ), + Cell( + why="TP4 eight-query boundary: sixteen CTA splits, masked short context", + spec=_cell(seq_len_q=8, kv_lens=(1500,), remap=True, layer_offset=4, holes=True), + ), + Cell( + why="TP4 nine-query boundary: eight CTA splits", + spec=_cell(seq_len_q=9, kv_lens=(4101,), remap=True, layer_offset=4), + ), + Cell( + why="TP4 twenty-two-query ceiling: eight CTA splits", + spec=_cell(seq_len_q=22, kv_lens=(4101,), remap=True, layer_offset=4), + ), + ) + + def __call__( + self, + q: torch.Tensor, + kv_pool: torch.Tensor, + indices: torch.Tensor, + seq_lens_kv: torch.Tensor, + mla_bmm1_scale: torch.Tensor, + mla_bmm2_scale: torch.Tensor, + output: torch.Tensor, + seq_len_q: int, + top_k: int, + block_table: Optional[torch.Tensor] = None, + tokens_per_block: int = 0, + page_index_scale: int = 0, + layer_offset: int = 0, + gather_before_wait: bool = False, + trigger_after_wait: bool = False, + ) -> torch.Tensor: + torch.ops.trtllm.glm_sparse_mla_flash_decode_latent( + q, + kv_pool, + indices, + seq_lens_kv, + mla_bmm1_scale, + mla_bmm2_scale, + output, + seq_len_q, + top_k, + block_table, + tokens_per_block, + page_index_scale, + layer_offset, + gather_before_wait, + trigger_after_wait, + ) + return output + + def reference( + self, + q: torch.Tensor, + kv_pool: torch.Tensor, + indices: torch.Tensor, + seq_lens_kv: torch.Tensor, + mla_bmm1_scale: torch.Tensor, + mla_bmm2_scale: torch.Tensor, + output: torch.Tensor, + seq_len_q: int, + top_k: int, + block_table: Optional[torch.Tensor] = None, + tokens_per_block: int = 0, + page_index_scale: int = 0, + layer_offset: int = 0, + gather_before_wait: bool = False, + trigger_after_wait: bool = False, + ) -> torch.Tensor: + """Selected-row attention with the kernel's split-local E4M3 probability rounding.""" + num_tokens = q.shape[0] + splits = 16 if num_tokens <= 8 else 8 + slots_per_split = MAX_TOP_K // splits + device = q.device + pool_rows = kv_pool.view(torch.uint8).reshape(-1, HEAD_DIM_QK) + + slots = torch.full((num_tokens, MAX_TOP_K), -1, dtype=torch.int64, device=device) + slots[:, :top_k] = indices[:, :top_k].long() + token = torch.arange(num_tokens, device=device) + seq = token // seq_len_q + num_valid = torch.clamp( + seq_lens_kv[seq].long() - (seq_len_q - 1 - token % seq_len_q), max=top_k + ) + if block_table is not None: + position = slots.clamp(min=0) + block = position // tokens_per_block + width = block_table.shape[1] + page = block_table[seq[:, None], block.clamp(max=width - 1)].long() + row = ( + page * (page_index_scale * tokens_per_block) + + position % tokens_per_block + + layer_offset * tokens_per_block + ) + resident = (slots >= 0) & (block < width) & (page >= 0) + else: + row = slots + resident = slots >= 0 + valid = resident & (torch.arange(MAX_TOP_K, device=device)[None, :] < num_valid[:, None]) + row = torch.where(valid, row, torch.zeros_like(row)) + + keys = pool_rows[row].view(torch.float8_e4m3fn).float() + queries = q.view(torch.float8_e4m3fn).float() + with true_fp32_matmul(): + scores = torch.einsum("thd,tjd->thj", queries, keys) * mla_bmm1_scale[1].float() + scores = scores.masked_fill(~valid[:, None, :], float("-inf")) + scores = scores.view(num_tokens, LATENT_NUM_HEADS, splits, slots_per_split) + split_max = scores.amax(dim=-1) + empty = split_max == float("-inf") + p = torch.exp2(scores - torch.where(empty, 0.0, split_max)[..., None]) + split_sum = p.sum(dim=-1) + p_e4m3 = (p * _P_SCALE).to(torch.float8_e4m3fn).float() + values = keys[..., :KV_LORA_RANK].reshape(num_tokens, splits, slots_per_split, KV_LORA_RANK) + with true_fp32_matmul(): + split_out = torch.einsum("thrj,trjd->thrd", p_e4m3, values) + merged_max = split_max.amax(dim=-1, keepdim=True) + alpha = torch.where(empty, 0.0, torch.exp2(split_max - merged_max)) + merged_sum = (alpha * split_sum).sum(dim=-1)[..., None] + merged = (alpha[..., None] * split_out).sum(dim=-2) + attn = torch.where( + merged_sum > 0, + merged / (merged_sum * _P_SCALE) * mla_bmm2_scale.reshape(-1)[0].float(), + 0.0, + ) + return attn.to(torch.bfloat16) + + def is_valid( + self, + q: torch.Tensor, + kv_pool: torch.Tensor, + indices: torch.Tensor, + seq_lens_kv: torch.Tensor, + mla_bmm1_scale: torch.Tensor, + mla_bmm2_scale: torch.Tensor, + output: torch.Tensor, + seq_len_q: int, + top_k: int, + block_table: Optional[torch.Tensor] = None, + tokens_per_block: int = 0, + page_index_scale: int = 0, + layer_offset: int = 0, + gather_before_wait: bool = False, + trigger_after_wait: bool = False, + ) -> None: + # Both ABIs read the same sparse metadata; this validator does not read v_b. + GlmSparseMlaFlashDecode.is_valid( + self, + q, + kv_pool, + indices, + seq_lens_kv, + mla_bmm1_scale, + mla_bmm2_scale, + None, + output, + seq_len_q, + top_k, + block_table, + tokens_per_block, + page_index_scale, + layer_offset, + gather_before_wait, + ) + + +sparse_mla_flash_decode_latent = GlmSparseMlaFlashDecodeLatent() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/_ipc_workspace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/_ipc_workspace.py new file mode 100644 index 000000000000..e0dba9607ef8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/_ipc_workspace.py @@ -0,0 +1,176 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Bind the engine's in-node IPC workspace and retain its actual allocations. + +The pointer table contains addresses, not payload. flags/layout are real mutable Tensors; +peer payload remains opaque cudaMalloc memory. This contract supports eager execution and +ordinary CUDA graphs, and deliberately does not support functionalization. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Sequence + +import torch + +from tensorrt_llm._ipc_utils import IpcMemory +from tensorrt_llm._torch.distributed import ops as distributed_ops +from tensorrt_llm._torch.distributed.allreduce_helper import force_all_reduce_deterministic +from tensorrt_llm._torch.distributed.communicator import Distributed +from tensorrt_llm.mapping import Mapping + + +@dataclass(frozen=True) +class IpcWorkspace: + table: torch.Tensor + flags: torch.Tensor + layout: torch.Tensor + owners: tuple[Any, ...] + rank: int + nranks: int + buffer_bytes: int + hidden: int + plain_rows: frozenset[int] + finalize_rows: frozenset[int] + + def args(self): + return self.table, self.flags, self.layout, self.rank, self.nranks, self.buffer_bytes + + +def _bind_local( + mapping: Mapping, workspaces: Sequence[torch.Tensor | None], max_rows: int, hidden: int +) -> IpcWorkspace: + if torch.compiler.is_compiling() or torch.cuda.is_current_stream_capturing(): + raise ValueError("IPC binding must happen eagerly before capture") + if ( + mapping.tp_size != 4 + or mapping.moe_ep_size != 1 + or mapping.moe_tp_size != 4 + or mapping.enable_attention_dp + or mapping.pp_size != 1 + or mapping.cp_size != 1 + or mapping.is_multi_node() + or hidden != 6144 + or not 1 <= max_rows <= 24 + ): + raise ValueError("IPC requires in-node TP4/EP1 BF16 hidden6144, at most24 rows") + if force_all_reduce_deterministic(): + raise ValueError("deterministic all-reduce policy uses the engine implementation") + if not workspaces or workspaces[0] is None: + raise ValueError("no existing all-reduce workspace") + table = workspaces[0] + if any(other is not table for other in workspaces): + raise ValueError("participating layers do not share the same IPC workspace") + # The public getter intentionally returns only the table. Read its ownership record once; + # do not call it here, since a cache miss would allocate/open an IPC group collectively. + cache = getattr(distributed_ops._thread_local, f"allreduce_workspaces_{mapping.pp_rank}", {}) + entry = cache.get(mapping) + if entry is None or len(entry) != 2 or entry[1] is not table: + raise ValueError("existing IPC ownership record is unavailable on this construction thread") + owners = tuple(entry[0]) + if len(owners) != 5 or not all(isinstance(x, IpcMemory) for x in owners[:3]): + raise ValueError("unexpected IPC owner layout") + if any( + x.mapping != mapping + or not x.open_ipc + or not x.local_ptr + or len(x.peer_ptrs) != 4 + or x.peer_ptrs[mapping.tp_rank] != x.local_ptr + or any(not ptr for ptr in x.peer_ptrs) + for x in owners[:3] + ): + raise ValueError("IPC peer allocations are not open") + flags, layout = owners[3:] + for name, tensor, dtype, count in ( + ("table", table, torch.int64, 14), + ("flags", flags, torch.int32, 3), + ("layout", layout, torch.int64, 2), + ): + if ( + not isinstance(tensor, torch.Tensor) + or not tensor.is_cuda + or tensor.device != table.device + or tensor.dtype != dtype + or tensor.ndim != 1 + or tensor.numel() != count + or not tensor.is_contiguous() + ): + raise ValueError(f"unexpected IPC {name} tensor") + expected = [ptr for owner in owners[:3] for ptr in owner.serialize()] + expected += [flags.data_ptr(), layout.data_ptr()] + # One initialization-time snapshot: never extract device state in forward/capture. + if table.cpu().tolist() != expected: + raise ValueError("IPC pointer table does not identify its retained owners") + buffer_bytes = int(layout.cpu().tolist()[1]) + if buffer_bytes <= 0 or buffer_bytes % 16 or max_rows * hidden * 2 * 4 > buffer_bytes: + raise ValueError("IPC Lamport buffer is too small for the planned shapes") + with torch.cuda.device(table.device): + plain = frozenset( + m + for m in range(1, max_rows + 1) + if torch.ops.trtllm.glm_intra_node_reduce_supported(m, hidden, 4, False) + ) + finalize = frozenset( + m + for m in range(1, max_rows + 1) + if torch.ops.trtllm.glm_intra_node_reduce_supported(m, hidden, 4, True) + ) + if not plain or not finalize: + raise ValueError("owned IPC kernels have no supported launch on this device") + return IpcWorkspace( + table, flags, layout, owners, mapping.tp_rank, 4, buffer_bytes, hidden, plain, finalize + ) + + +def bind_ipc_workspace( + mapping: Mapping, + workspaces: Sequence[torch.Tensor | None], + *, + max_rows: int = 24, + hidden: int = 6144, +) -> tuple[IpcWorkspace | None, str]: + """Called on every rank in the same order, once after ordinary layer construction. + + A rank-local failure becomes a uniform unsupported result. There is no allocation, + cache mutation, flag reset, or repeated lookup in model forward. + """ + binding = None + reason = "" + try: + binding = _bind_local(mapping, workspaces, max_rows, hidden) + except Exception as exc: + reason = f"{type(exc).__name__}: {exc}" + status = ( + reason, + sorted(binding.plain_rows) if binding else [], + sorted(binding.finalize_rows) if binding else [], + ) + results = Distributed.get(mapping).tp_allgather(status, small_payload=True) + failures = [f"rank {rank}: {item[0]}" for rank, item in enumerate(results) if item[0]] + if failures: + return None, "; ".join(failures) + if any(item[1:] != results[0][1:] for item in results): + return None, "IPC launch capabilities differ between ranks" + return binding, "" + + +def validate_ipc_operands(residual, gamma, table, flags, layout, rank, nranks, buffer_bytes): + """Executable entry contract, called by tests through validating(), never by forward.""" + assert residual.ndim == 2 and 1 <= residual.shape[0] <= 24 and residual.shape[1] == 6144 + assert residual.dtype == gamma.dtype == torch.bfloat16 and gamma.shape == (6144,) + assert nranks == 4 and 0 <= rank < nranks + assert table.shape == (14,) and table.dtype == torch.int64 + assert flags.shape == (3,) and flags.dtype == torch.int32 + assert layout.shape == (2,) and layout.dtype == torch.int64 + for value in (residual, gamma, table, flags, layout): + assert value.is_cuda and value.device == residual.device and value.is_contiguous() + assert buffer_bytes >= residual.numel() * 2 * nranks and buffer_bytes % 16 == 0 + assert table.cpu().tolist()[-2:] == [flags.data_ptr(), layout.data_ptr()] + assert layout[1].item() == buffer_bytes + + +def residual_norm_reference(total, residual, gamma, epsilon): + summed = total.to(torch.bfloat16) + residual + normed = summed.float() * torch.rsqrt(summed.float().square().mean(-1, keepdim=True) + epsilon) + return (normed * gamma.float()).to(torch.bfloat16), summed diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/_mnnvl_workspace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/_mnnvl_workspace.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/_mnnvl_workspace.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/_mnnvl_workspace.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/intra_node_moe_finalize.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/intra_node_moe_finalize.py new file mode 100644 index 000000000000..3e613d6422c9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/intra_node_moe_finalize.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Local top-k FP32 finalize followed by rank-order BF16 IPC sum and residual/RMSNorm.""" + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._ipc_workspace import residual_norm_reference, validate_ipc_operands + + +class GlmIntraNodeMoeFinalize(OpWrapper): + ARCHS = frozenset({Arch.SM_107}) + CELLS = tuple( + Cell( + why=f"TP4 routed verify rows={rows}", + spec=dict(rows=rows, hidden=6144, world=4, top_k=8), + ) + for rows in (1, 3, 6, 12, 18) + ) + note = ( + "Uses real IPC flags/layout and preserves BF16 rank addition; no functionalization support." + ) + + def __call__( + self, + fc2, + weights, + indices, + shared, + residual, + gamma, + epsilon, + table, + flags, + layout, + rank, + nranks, + buffer_bytes, + ): + return torch.ops.trtllm.glm_intra_node_moe_finalize( + fc2, + weights, + indices, + shared, + residual, + gamma, + epsilon, + table, + flags, + layout, + rank, + nranks, + buffer_bytes, + ) + + def reference( + self, + per_rank_fc2, + per_rank_weights, + per_rank_indices, + per_rank_shared, + residual, + gamma, + epsilon, + ): + local = [] + for fc2, weights, indices, shared in zip( + per_rank_fc2, per_rank_weights, per_rank_indices, per_rank_shared + ): + acc = torch.zeros_like(residual, dtype=torch.float32) + for slot in range(indices.shape[1]): + ids = indices[:, slot].long() + rows = fc2[ids.clamp_min(0)].float() + rows = torch.where((ids >= 0).unsqueeze(-1), rows, 0) + acc = acc + rows * weights[:, slot].float().unsqueeze(-1) + local.append((acc + shared.float()).to(torch.bfloat16)) + total = local[0] + for value in local[1:]: + total = total + value + return residual_norm_reference(total, residual, gamma, epsilon) + + def is_valid( + self, + fc2, + weights, + indices, + shared, + residual, + gamma, + epsilon, + table, + flags, + layout, + rank, + nranks, + buffer_bytes, + ): + validate_ipc_operands(residual, gamma, table, flags, layout, rank, nranks, buffer_bytes) + assert fc2.ndim == 2 and fc2.shape[1] == residual.shape[1] + assert indices.ndim == 2 and indices.shape[0] == residual.shape[0] + assert 1 <= indices.shape[1] <= 8 and indices.dtype == torch.int32 + assert weights.shape == indices.shape and weights.dtype == torch.bfloat16 + assert shared.shape == residual.shape and shared.dtype == fc2.dtype == torch.bfloat16 + for value in (fc2, weights, indices, shared): + assert value.is_contiguous() and value.device == residual.device + assert ((indices >= -1) & (indices < fc2.shape[0])).all().item() + + +intra_node_moe_finalize = GlmIntraNodeMoeFinalize() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/intra_node_reduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/intra_node_reduce.py new file mode 100644 index 000000000000..a95db214a3cd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/intra_node_reduce.py @@ -0,0 +1,182 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""TP4 residual/RMSNorm, MoE finalize and LINEAR FP4 on one retained IPC workspace. + +All calls on an instance must be stream ordered, including engine collective fallbacks. +The shared workspace's prior dirty range may be larger than the current rows. The native +kernels clear that complete range and preserve the shared rotation protocol. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Sequence + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm.mapping import Mapping + +from ._ipc_workspace import ( + IpcWorkspace, + bind_ipc_workspace, + residual_norm_reference, + validate_ipc_operands, +) +from .intra_node_moe_finalize import intra_node_moe_finalize +from .pre_moe_fp4_linear import pre_moe_fp4_linear + + +class GlmIntraNodeReduce(OpWrapper): + """BF16 residual/RMSNorm over rank-ordered partials on the bound TP4 IPC workspace.""" + + ARCHS = frozenset({Arch.SM_107}) + CELLS = tuple( + Cell(why=f"TP4 verify/draft rows={rows}", spec=dict(rows=rows, hidden=6144, world=4)) + for rows in (1, 2, 3, 6, 12, 18) + ) + note = "Opaque peer payload supports eager/CUDA graph only; actual flags/layout are mutated." + + def __call__( + self, + input, + residual, + gamma, + epsilon, + table, + flags, + layout, + rank, + nranks, + buffer_bytes, + trigger_completion_at_end, + ): + return torch.ops.trtllm.glm_intra_node_reduce( + input, + residual, + gamma, + epsilon, + table, + flags, + layout, + rank, + nranks, + buffer_bytes, + trigger_completion_at_end, + ) + + def reference(self, per_rank_inputs, residual, gamma, epsilon, *, fp32_acc=True): + total = per_rank_inputs[0].float() + for value in per_rank_inputs[1:]: + total = ( + total + value.float() if fp32_acc else (total.to(torch.bfloat16) + value).float() + ) + return residual_norm_reference(total, residual, gamma, epsilon) + + def is_valid( + self, + input, + residual, + gamma, + epsilon, + table, + flags, + layout, + rank, + nranks, + buffer_bytes, + trigger_completion_at_end, + ): + validate_ipc_operands(residual, gamma, table, flags, layout, rank, nranks, buffer_bytes) + assert input.shape == residual.shape and input.dtype == residual.dtype + assert input.is_contiguous() and input.device == residual.device + assert input.data_ptr() % 16 == residual.data_ptr() % 16 == gamma.data_ptr() % 16 == 0 + + +intra_node_reduce = GlmIntraNodeReduce() + + +@dataclass(frozen=True) +class IpcCollectives: + workspace: IpcWorkspace + + @classmethod + def bind( + cls, + mapping: Mapping, + workspaces: Sequence[torch.Tensor | None], + *, + max_rows: int = 24, + hidden: int = 6144, + ) -> tuple[IpcCollectives | None, str]: + workspace, reason = bind_ipc_workspace( + mapping, workspaces, max_rows=max_rows, hidden=hidden + ) + return (cls(workspace) if workspace is not None else None), reason + + def supported(self, rows: int, *, finalize: bool = False) -> bool: + if torch.compiler.is_compiling(): + return False + allowed = self.workspace.finalize_rows if finalize else self.workspace.plain_rows + return rows in allowed + + def _check(self, rows: int, *, finalize: bool = False): + if not self.supported(rows, finalize=finalize): + raise ValueError( + "unsupported owned IPC call; select the engine path before writing state" + ) + + def plain_reduce( + self, + norm, + partial: torch.Tensor, + residual: torch.Tensor, + *, + trigger_completion_at_end: bool = True, + ): + """FP32 rank sum (unless explicitly opted out), BF16 residual add, RMSNorm.""" + self._check(partial.shape[0]) + return intra_node_reduce( + partial, + residual, + norm.weight, + norm.variance_epsilon, + *self.workspace.args(), + trigger_completion_at_end, + ) + + def finalize( + self, + fc2: torch.Tensor, + expert_weights: torch.Tensor, + permuted_rows: torch.Tensor, + shared: torch.Tensor, + residual: torch.Tensor, + norm, + ): + """Local top-k in FP32, one BF16 rounding, rank-order BF16 sum, residual/RMSNorm.""" + self._check(residual.shape[0], finalize=True) + return intra_node_moe_finalize( + fc2, + expert_weights, + permuted_rows, + shared, + residual, + norm.weight, + norm.variance_epsilon, + *self.workspace.args(), + ) + + def premoe(self, partial: torch.Tensor, residual: torch.Tensor, norm, scale: torch.Tensor): + """Return norm/residual plus packed FP4 and LINEAR SF using FC31 input scale. + + ``scale`` is backend.fc31_input_scale: one CUDA float32 value. Packed output is + uint8 [M,3072], scale factors uint8 [M,384]. The consumer must bypass requantization. + ``residual`` must be ready before the producer's PDL trigger: this kernel reads it + before its dependency wait. Newly produced residuals need an ordinary stream dependency. + """ + self._check(partial.shape[0]) + return pre_moe_fp4_linear( + partial, residual, norm.weight, norm.variance_epsilon, scale, *self.workspace.args() + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_allgather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/mnnvl_allgather.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_allgather.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/mnnvl_allgather.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_hop.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/mnnvl_hop.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_hop.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/mnnvl_hop.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_moe_finalize_hop.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/mnnvl_moe_finalize_hop.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/comm/mnnvl_moe_finalize_hop.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/mnnvl_moe_finalize_hop.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/pre_moe_fp4_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/pre_moe_fp4_linear.py new file mode 100644 index 000000000000..8b4ca42b5b59 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/comm/pre_moe_fp4_linear.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared native residual/RMSNorm/FP4 kernel with an owned LINEAR-scale consumer contract.""" + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ._ipc_workspace import residual_norm_reference, validate_ipc_operands + + +class GlmPreMoeFp4Linear(OpWrapper): + ARCHS = frozenset({Arch.SM_107}) + CELLS = tuple( + Cell(why=f"TP4 target pre-MoE FP4 rows={rows}", spec=dict(rows=rows, hidden=6144, world=4)) + for rows in (1, 3, 18) + ) + note = "scale is FC31 input scale; the bound table indirectly writes the real flags/layout operands." + + def __call__( + self, + input, + residual, + gamma, + epsilon, + scale, + table, + flags, + layout, + rank, + nranks, + buffer_bytes, + ): + return torch.ops.trtllm.glm_pre_moe_fp4_linear( + input, residual, gamma, epsilon, scale, table, flags, layout, rank, nranks, buffer_bytes + ) + + def reference(self, per_rank_inputs, residual, gamma, epsilon, scale, *, fp32_acc=True): + total = per_rank_inputs[0].float() + for value in per_rank_inputs[1:]: + total = ( + total + value.float() if fp32_acc else (total.to(torch.bfloat16) + value).float() + ) + norm, res = residual_norm_reference(total, residual, gamma, epsilon) + # Quantization's shared plain-Torch oracle, not the native op under test. + from ..quantization.fp4_quantize import fp4_quantize + + packed, sf = fp4_quantize.reference(norm, scale, 16, False, False) + return norm, res, packed, sf + + def is_valid( + self, + input, + residual, + gamma, + epsilon, + scale, + table, + flags, + layout, + rank, + nranks, + buffer_bytes, + ): + validate_ipc_operands(residual, gamma, table, flags, layout, rank, nranks, buffer_bytes) + assert input.shape == residual.shape and input.dtype == residual.dtype + assert input.is_contiguous() and input.device == residual.device + assert scale.dtype == torch.float32 and scale.numel() == 1 and scale.is_contiguous() + assert scale.device == residual.device + + +pre_moe_fp4_linear = GlmPreMoeFp4Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_fp8_paged_mqa_logits.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/_fp8_paged_mqa_logits.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_fp8_paged_mqa_logits.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/_fp8_paged_mqa_logits.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_fp8_paged_mqa_logits_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/_fp8_paged_mqa_logits_kernel.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_fp8_paged_mqa_logits_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/_fp8_paged_mqa_logits_kernel.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/_full_context_topk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/_full_context_topk.py new file mode 100644 index 000000000000..f734b04b61d8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/_full_context_topk.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Capture-stable K=2048 streaming TopK, sharing the SMs with MLA q_b. + +The device implementation and its per-row causal lengths are the common GVR kernels. +Only the small-batch streaming launch policy belongs to this model. Workspace is +caller-owned and must not be shared by concurrent launches. +""" + +from __future__ import annotations + +import functools + +import torch + +from tensorrt_llm._torch.kernels.blackwell.top_k import gvr_topk_decode_self_sampling as device +from tensorrt_llm._torch.kernels.blackwell.top_k.gvr_topk_decode_self_sampling_host import ( + kernel_view, + validate_run_ws, + workspace_bytes, +) + +TOPK = 2048 +MAX_ROWS = 24 + + +def streaming_plan(rows: int, envelope: int, num_sms: int, sm_reserve: int) -> tuple: + """Return the current streaming-main ABI with a cap applied before deriving its fields.""" + if ( + not 1 <= rows <= min(MAX_ROWS, num_sms - sm_reserve) + or not 0 <= sm_reserve < num_sms + or envelope < 1 + ): + raise ValueError("unsupported small-batch TopK launch geometry") + n_route = max(envelope, TOPK + 1) + splits = min( + max(148 // rows, 1), + max(((n_route >> 2) + 1023) // 1024, 1), + max((num_sms - sm_reserve) // rows, 1), + ) + # rows*splits <=148: this is always the large-CTA streaming branch. + scap = 16384 if splits == 1 else 8192 + cmp = 4096 + q = ((n_route >> 2) + splits - 1) // splits + per = q >> 10 + unroll = 8 if per >= 8 else 4 if per >= 4 else 2 if per >= 2 else 1 + tpl = (1024, unroll, 1, 256, 2, splits > 1, False) + aim = 4 * TOPK if splits == 1 else 2 * TOPK + sfac = 64 if splits == 1 else 32 if splits == 2 else 48 + amin = 3 * TOPK if splits == 2 else 7 * TOPK // 2 + return tpl, splits, scap, cmp, (aim, sfac, amin, 0, 0) + + +@functools.cache +def _device_profile(index: int) -> tuple[int, int]: + props = torch.cuda.get_device_properties(index) + return props.multi_processor_count, 10 * props.major + props.minor + + +_LAUNCHERS = {} + + +def _launcher(rows, stride, envelope, next_n, num_sms, sm_version, sm_reserve, deterministic): + key = (rows, stride, envelope, next_n, num_sms, sm_version, sm_reserve, deterministic) + hit = _LAUNCHERS.get(key) + if hit is not None: + return hit + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("warm up this full-context TopK shape before graph capture") + tpl, splits, scap, cmp, tail = streaming_plan(rows, envelope, num_sms, sm_reserve) + fn = device.get_compiled( + tpl + (next_n, 0, splits), + hint_free=True, + det=deterministic, + fused_k=TOPK, + warp_rank=sm_version in (100, 103), + staged_k=TOPK if splits == 1 and sm_version in (100, 103) else 0, + ) + # The varlen kernel reads the dynamic causal lengths from kv_lens. The first + # scalar n and the five sampling-ladder scalars are dead ABI slots in this mode. + pre = (0, stride, TOPK, scap, cmp, splits, 0, 0, 0, 0, 0) + result = fn, pre, tail + _LAUNCHERS[key] = result + return result + + +def full_context_topk( + logits: torch.Tensor, + kv_lens: torch.Tensor, + indices: torch.Tensor, + workspace: torch.Tensor, + *, + next_n: int, + max_seq_len: int, + sm_reserve: int, + deterministic: bool, +) -> torch.Tensor: + """Write int32 indices[rows,2048] from FP32 logits and device request lengths. + + Each request has ``next_n`` rows; row i sees ``kv_len-next_n+i+1`` keys. + Short rows produce identity indices followed by -1. The physical row pitch + and context envelope are capture-stable; request lengths remain device data. + """ + if logits.ndim != 2 or logits.dtype != torch.float32 or not logits.is_cuda: + raise ValueError("logits must be a CUDA FP32 matrix") + rows = logits.shape[0] + if not 1 <= rows <= MAX_ROWS or next_n < 1 or rows % next_n: + raise ValueError("unsupported number of TopK rows or query length") + if ( + indices.shape != (rows, TOPK) + or indices.dtype != torch.int32 + or not indices.is_contiguous() + or indices.device != logits.device + ): + raise ValueError("indices must be contiguous int32[rows,2048] on the logits device") + if ( + kv_lens.shape != (rows // next_n,) + or kv_lens.dtype != torch.int32 + or not kv_lens.is_contiguous() + or kv_lens.device != logits.device + ): + raise ValueError("kv_lens must contain one contiguous CUDA int32 length per request") + stride = logits.stride(0) + if ( + logits.stride(1) != 1 + or stride < logits.shape[1] + or stride % 4 + or logits.data_ptr() % 16 + or not 1 <= max_seq_len <= logits.shape[1] + ): + raise ValueError("logits must have an aligned row-major physical envelope") + if logits.untyped_storage().nbytes() // 4 < logits.storage_offset() + rows * stride: + raise ValueError("logits storage does not cover its physical row pitch") + validate_run_ws(workspace, logits) + if not workspace.is_contiguous(): + raise ValueError("TopK workspace must be contiguous") + lg = logits.as_strided((rows, stride), (stride, 1)) + num_sms, sm_version = _device_profile(logits.device.index) + fn, pre, tail = _launcher( + rows, + stride, + min(max_seq_len, stride), + next_n, + num_sms, + sm_version, + sm_reserve, + bool(deterministic), + ) + # hint_free=True makes the pre_idx ABI slot unread; no separate hint buffer is needed. + fn(lg, indices, indices, kernel_view(workspace), *pre, kv_lens, *tail) + return indices + + +__all__ = ["MAX_ROWS", "TOPK", "full_context_topk", "streaming_plan", "workspace_bytes"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_hist_topk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/_hist_topk.py similarity index 99% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_hist_topk.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/_hist_topk.py index a9973f360be2..9603a54c090b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_hist_topk.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/_hist_topk.py @@ -7,7 +7,7 @@ scores only its contiguous shard of the request's indexer keys (``glm_dsa_hist_shard_meta`` builds the shard's context length and block table), and ``glm_dsa_hist_topk`` finds every row's exact top-2048 over all shards by exchanging key histograms through a private multicast workspace; every -rank ends with the same rows. See ``cpp/tensorrt_llm/kernels/glmMoeDsa/glmDsaHistTopK.h``. +rank ends with the same rows. See ``cpp/tensorrt_llm/kernels/glmMoeDsaRubin/glmDsaHistTopK.h``. """ from __future__ import annotations diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_indexer_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/_indexer_prologue.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_indexer_prologue.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/_indexer_prologue.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_indexer_reference.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/_indexer_reference.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_indexer_reference.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/_indexer_reference.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_mtp_topk_stash.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/_mtp_topk_stash.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/_mtp_topk_stash.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/_mtp_topk_stash.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/convert_req_index_to_global.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/convert_req_index_to_global.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/convert_req_index_to_global.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/convert_req_index_to_global.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/fp8_paged_mqa_logits.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/fp8_paged_mqa_logits.py similarity index 94% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/fp8_paged_mqa_logits.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/fp8_paged_mqa_logits.py index 6ef531814cdb..fbb933ac388f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/fp8_paged_mqa_logits.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/fp8_paged_mqa_logits.py @@ -89,6 +89,18 @@ class GlmFp8PagedMqaLogits(OpWrapper): why="four requests of 8 rows each (a larger batch's verify step)", spec=_spec([3_000, 70_000, 129, 40_000], max_context_len=131_072), ), + Cell( + why="TP4 three-request verify, full six-query atom and unequal contexts", + spec=_spec([3000, 33000, 129], max_context_len=65536, next_n=6), + ), + Cell( + why="TP4 partial one-request verify, full six-query atom", + spec=_spec([33000], max_context_len=65536, next_n=6), + ), + Cell( + why="TP4 three-request draft, one-query atoms", + spec=_spec([3000, 33000, 129], max_context_len=65536, next_n=1, early_kv=False), + ), ) note = """ @@ -107,7 +119,7 @@ class GlmFp8PagedMqaLogits(OpWrapper): kernels it waits for write no key older than each request's newest ``next_n`` and no block table, context length or schedule entry. - Not certified: FP16 epilogues, next_n other than 8, head counts other than 32, pages other + Not certified: FP16 epilogues, next_n other than 1, 6 or 8, head counts other than 32, pages other than 64 tokens. """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/hist_shard_meta.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/hist_shard_meta.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/hist_shard_meta.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/hist_shard_meta.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/hist_topk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/hist_topk.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/hist_topk.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/hist_topk.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/indexer_k_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/indexer_k_tail.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/indexer_k_tail.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/indexer_k_tail.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/indexer_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/indexer_prologue.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/indexer_prologue.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/indexer_prologue.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/paged_mqa_logits_schedule.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/paged_mqa_logits_schedule.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/paged_mqa_logits_schedule.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/paged_mqa_logits_schedule.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/stash_last_accepted_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/stash_last_accepted_rows.py similarity index 89% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/stash_last_accepted_rows.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/stash_last_accepted_rows.py index 02b6d14a1300..2fe4048dd977 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/dsa/stash_last_accepted_rows.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/dsa/stash_last_accepted_rows.py @@ -29,6 +29,14 @@ class GlmStashLastAcceptedRows(OpWrapper): ARCHS = frozenset({Arch.SM_107}) CELLS: tuple[Cell, ...] = ( + Cell( + why="TP4 first draft: three requests with mixed acceptance", + spec=_spec(3, [1, 4, 6], next_n=6), + ), + Cell( + why="TP8 first draft: two requests with mixed acceptance", + spec=_spec(2, [1, 4], next_n=4, dst="wide"), + ), Cell(why="c1: one request, 8 TopK rows, 3 tokens accepted", spec=_spec(1, [3])), Cell(why="c1: every draft token accepted (the last row)", spec=_spec(1, [8])), Cell(why="counts past the window clamp: 0 and 9", spec=_spec(2, [0, 9])), @@ -40,7 +48,7 @@ class GlmStashLastAcceptedRows(OpWrapper): note = """ Metadata consumed: none. The target passes the draft step-0 TopK output - ([num_requests * 8, 2048] int32), the verify step's accepted counts per request (int32 on the + ([num_requests * next_n, 2048] int32), the verify step's accepted counts per request (int32 on the device) and the rows [0, num_requests) of the persistent shared-TopK buffer. A plain launch (no PDL): at c1 it runs on its own stream behind an event recorded after the TopK, and the layer joins it at the end of its attention. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/embedding/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/embedding/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/embedding/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/embedding/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/embedding/fused_embedding.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/embedding/fused_embedding.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/embedding/fused_embedding.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/embedding/fused_embedding.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/_low_m_splitk.py similarity index 90% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/_low_m_splitk.py index 8c592c735b35..3e6cf6fc9e02 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/_low_m_splitk.py @@ -74,6 +74,9 @@ #: Size and alignment of the TMEM base pointer slot. _TMEM_POINTER_BYTES = 4 +#: Default dense GEMM tensor-memory allocation; small batched v_b can use fewer columns. +_TMEM_COLS = 256 + #: Bytes per BF16/FP16 element. _AB_ELEMENT_BYTES = 2 @@ -388,6 +391,42 @@ def splitk_default_tactic( ] +def one_wave_tactic( + tactic: SplitKTactic, + m: int, + n: int, + k: int, + *, + sm_count: int, + smem_capacity: int = _SMEM_CAPACITY_BYTES, +) -> SplitKTactic: + """``tactic`` with 128-row instead of 64-row tiles when its CTAs (one per SM) need more than one wave and the + 128-row tiles fit in one, else ``tactic`` unchanged. + + Each output keeps its split-K slice and its K order, so the results are the same; the grid is half as large and + has no partial second wave (the c1 LM head: 152 instead of 303 CTAs). Only long-K tactics (more than four K + tiles, the full-depth pipelines) are refined. + """ + + def cta_count(t: SplitKTactic) -> int: + return ((n + t.mma_m - 1) // t.mma_m) * ((m + t.mma_n - 1) // t.mma_n) * t.split_k + + if tactic.mma_m != 64 or k <= 4 * _CTA_K or cta_count(tactic) <= sm_count: + return tactic + candidate = dataclasses.replace(tactic, mma_m=128) + if cta_count(candidate) > sm_count: + return tactic + stages = min(tactic.ab_stages, _max_ab_stages_for(candidate, smem_capacity)) + if stages < _DEFAULT_MIN_AB_STAGES: + return tactic + candidate = dataclasses.replace(candidate, ab_stages=stages) + try: + validate_tactic(candidate, m, n, k) + except ValueError: + return tactic + return candidate + + @dsl_user_op def _map_shared_rank( smem_ptr: "cute.Pointer", @@ -667,6 +706,7 @@ def __init__( reserve_full_smem: bool = True, publish: LamportPublish | None = None, hnorm_b: HnormB | None = None, + tmem_cols: int = _TMEM_COLS, ) -> None: self.acc_dtype = cutlass.Float32 self.cta_m = tactic.mma_m @@ -679,6 +719,14 @@ def __init__( # Shared memory requested at launch: the full capacity (one CTA per SM) or only the # tactic's footprint. The kernel itself is the same either way. self.reserve_full_smem = reserve_full_smem + if ( + tmem_cols < 32 + or tmem_cols > 512 + or tmem_cols & (tmem_cols - 1) + or tmem_cols < tactic.mma_n + ): + raise ValueError(f"tmem_cols {tmem_cols} does not hold the {tactic.mma_n}-column tile") + self.tmem_cols = tmem_cols self.launch_smem_bytes = tactic_launch_smem_bytes(tactic) # Lamport publish (run_splitk_publish): the owner stages the reduced FP32 tile in shared # memory and stores it into the all-reduce workspace instead of writing C. @@ -1141,7 +1189,7 @@ def mma_warp( mma_inst_tile_k: cutlass.Constexpr, k_tile_count: cutlass.Int32, ): - num_tmem_cols = 256 + num_tmem_cols = self.tmem_cols cute.arch.alloc_tmem(num_tmem_cols, tmem_base_ptr, is_two_cta=False) cute.arch.mbarrier_arrive(bar_tmem_alloc) cute.arch.relinquish_tmem_alloc_permit(is_two_cta=False) @@ -1556,6 +1604,7 @@ def _get_compiled_splitk_kernel( leading_dims: tuple[int, int, int], reserve_full_smem: bool = True, publish: LamportPublish | None = None, + tmem_cols: int = _TMEM_COLS, ): if dtype not in _SUPPORTED_TORCH_DTYPES: raise ValueError(f"split-K dense GEMM supports {_SUPPORTED_TORCH_DTYPES}; got {dtype}") @@ -1566,6 +1615,7 @@ def _get_compiled_splitk_kernel( has_bias=has_bias, reserve_full_smem=reserve_full_smem, publish=publish, + tmem_cols=tmem_cols, ) compile_tensors = _make_compile_repr_tensors(dtype, has_bias, *leading_dims) stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) @@ -1683,6 +1733,136 @@ def run_splitk_dense( return out +BMM_MAX_M = 64 +_BMM_PREWAIT_MMA_M = 64 + + +def _validate_bmm_tensors(a, w, out) -> tuple[int, int, int, int]: + tensors = (a, w, out) + if any(not isinstance(tensor, _torch.Tensor) for tensor in tensors): + raise ValueError("a, w, and out must be torch tensors") + if a.ndim != 3 or w.ndim != 3 or out.ndim != 3: + raise ValueError("the batched split-K GEMM accepts only 3D tensors") + if a.device.type != "cuda" or any(tensor.device != a.device for tensor in tensors): + raise ValueError("all tensors must be on the same CUDA device") + if a.dtype not in _SUPPORTED_TORCH_DTYPES or any(tensor.dtype != a.dtype for tensor in tensors): + raise ValueError("a, w, and out must share BF16 or FP16 dtype") + batch, m, k = a.shape + if w.shape[0] != batch or w.shape[2] != k: + raise ValueError(f"incompatible shapes: a is {tuple(a.shape)}, w is {tuple(w.shape)}") + n = w.shape[1] + if tuple(out.shape) != (batch, m, n): + raise ValueError(f"out must have shape {(batch, m, n)}, got {tuple(out.shape)}") + if not w.is_contiguous() or a.stride(2) != 1 or out.stride(2) != 1: + raise ValueError("w must be contiguous, a and out contiguous in their last dimension") + # TMA loads a and w: their outer byte strides must be multiples of 16 bytes. + for tensor, tname in ((a, "a"), (w, "w")): + if any((stride * tensor.element_size()) % 16 for stride in tensor.stride()[:2]): + raise ValueError( + f"tensor '{tname}' strides {tensor.stride()} are not 16-byte multiples" + ) + if any(tensor.data_ptr() % 32 for tensor in tensors): + raise ValueError("a, w, and out must be 32-byte aligned") + if batch < 1 or m < 1 or n < 1 or k < 1: + raise ValueError("batched GEMM dimensions must be positive") + if any(operands_overlap(out, operand) for operand in (a, w)): + raise ValueError("out must not overlap either input") + # Transposed [M, heads, width] views are supported, overlapping rows are not. + for tensor in (a, out): + extent = 1 + for size, stride in sorted(zip(tensor.shape, tensor.stride()), key=lambda pair: pair[1]): + if size > 1 and stride < extent: + raise ValueError("batched GEMM rows must not overlap") + extent += (size - 1) * stride + return batch, m, n, k + + +def bmm_prewait_tactic(m: int, n: int, k: int) -> SplitKTactic | None: + """The tactic of a batched GEMM whose weights load before its grid-dependency wait: 64-row + kernel-M tiles, the smallest kernel-N tile of 8 / 16 / 32 rows that holds the rows (32-row + tiles above 32), no split-K, and one stage per K tile, so that every CTA requests its whole + weight tile before the wait and sums K in order. None where no such tactic is valid.""" + rows = min(m, MAX_M) + mma_n = 8 if rows <= 8 else 16 if rows <= 16 else 32 + if k <= 0 or k % _CTA_K: + return None + stages = k // _CTA_K + if stages > _MAX_AB_STAGES: + return None + tactic = SplitKTactic(_BMM_PREWAIT_MMA_M, mma_n, 1, stages) + try: + validate_tactic(tactic, rows, n, k) + except ValueError: + return None + return tactic + + +def bmm_supported(a: _torch.Tensor, w: _torch.Tensor, out: _torch.Tensor) -> bool: + """Whether ``run_splitk_bmm`` serves these tensors (shapes, layout, alignment, row count).""" + try: + _, m, _, _ = _validate_bmm_tensors(a, w, out) + except ValueError: + return False + return 1 <= m <= BMM_MAX_M + + +def bmm_tmem_cols(tactic: SplitKTactic) -> int: + """Tensor-memory columns that hold the FP32 accumulator of ``tactic`` (a power of two of at + least 32), so that CTAs of small tiles that share an SM do not wait for each other's + allocation.""" + cols = 32 + while cols < tactic.mma_n: + cols *= 2 + return cols + + +def run_splitk_bmm( + a: _torch.Tensor, + w: _torch.Tensor, + out: _torch.Tensor, + pdl: bool, + tactic: SplitKTactic, + reserve_full_smem: bool = True, + tmem_cols: int | None = None, +) -> _torch.Tensor: + """Run ``out[l] = a[l] @ w[l]^T`` for every batch ``l`` in one launch of the split-K kernel. + + Args: + a: ``[L, M, K]`` BF16/FP16 activations, contiguous in K (any M and L strides). + w: ``[L, N, K]`` contiguous weights (the kernel's A operand: loaded before the grid + dependency wait). + out: ``[L, M, N]`` output, contiguous in N (any M and L strides). + pdl, tactic, reserve_full_smem: as for ``run_splitk_dense``. M up to ``BMM_MAX_M``; + more rows than ``tactic.mma_n`` take more kernel-N tiles. + tmem_cols: tensor-memory columns per CTA; None allocates ``bmm_tmem_cols(tactic)``. + + The grid's third dimension walks the batch: each CTA computes one ``(mma_m, mma_n)`` tile + of one batch over its K slice, with the accumulation of ``run_splitk_dense``. + """ + batch, m, n, k = _validate_bmm_tensors(a, w, out) + if not 1 <= m <= BMM_MAX_M: + raise ValueError(f"the batched split-K GEMM requires 1 <= M <= {BMM_MAX_M}, got {m}") + validate_tactic(tactic, min(m, MAX_M), n, k) + swapped = (w, a.transpose(-2, -1), out.transpose(-2, -1)) + leading_dims = tuple(_detect_leading_dim(tensor) for tensor in swapped) + cute_tensors = tuple( + _from_dlpack_dynamic(tensor, leading_dim) + for tensor, leading_dim in zip(swapped, leading_dims, strict=True) + ) + compiled = _get_compiled_splitk_kernel( + dtype=a.dtype, + tactic=tactic, + use_pdl=pdl, + has_bias=False, + leading_dims=leading_dims, + reserve_full_smem=reserve_full_smem, + tmem_cols=bmm_tmem_cols(tactic) if tmem_cols is None else tmem_cols, + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + compiled(*cute_tensors, stream) + return out + + def publish_tactic( tactic: SplitKTactic, *, smem_capacity: int = _SMEM_CAPACITY_BYTES ) -> SplitKTactic: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk_fused.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/_low_m_splitk_fused.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/_low_m_splitk_fused.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/_low_m_splitk_fused.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/dense_mlp_small_m.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/dense_mlp_small_m.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/dense_mlp_small_m.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/dense_mlp_small_m.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/draft_hidden_half.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/draft_hidden_half.py similarity index 74% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/draft_hidden_half.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/draft_hidden_half.py index 99bdabd41f32..4fa308aa23ed 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/draft_hidden_half.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/draft_hidden_half.py @@ -16,10 +16,13 @@ from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl from ._low_m_splitk import ( + SplitKTactic, device_limits, + hnorm_b_supported, operands_overlap, run_splitk_hidden_half, splitk_default_tactic, + validate_tactic, ) #: GLM-5.2 at tp16: hidden 6144, this rank's 6144 / 16 = 384 output rows of eh_proj's hidden @@ -27,6 +30,29 @@ _GLM = dict(k=6144, n=384, eps=1e-5) +def hidden_half_tactic( + rows: int, + n: int, + k: int, + *, + in_kernel_norm: bool, + sm_count: int, + smem_capacity: int, +) -> SplitKTactic: + """Select a GEMM tile whose stages also hold the fused norm's complete K slice.""" + tactic = splitk_default_tactic(rows, n, k, sm_count=sm_count, smem_capacity=smem_capacity) + if in_kernel_norm: + if rows != 1: + raise ValueError("the in-kernel norm needs one activation row") + if (rows, n, k) == (1, 768, 6144) and not hnorm_b_supported(tactic, k): + # A 64-row weight tile leaves room for all six K tiles of the norm producer. + tactic = SplitKTactic(mma_m=64, mma_n=8, split_k=8, ab_stages=6) + validate_tactic(tactic, rows, n, k, smem_capacity=smem_capacity) + if tactic.split_k > sm_count or not hnorm_b_supported(tactic, k): + raise ValueError(f"the in-kernel norm cannot serve {(rows, n, k)}") + return tactic + + class GlmDraftHiddenHalf(OpWrapper): """`out = rmsnorm(x) @ weight.T` (with `norm_weight`) or `out = x @ weight.T` (without), FP32 out; returns `out`. @@ -56,11 +82,20 @@ class GlmDraftHiddenHalf(OpWrapper): why="two requests per draft step: 2 normalized rows, the plain FP32-output GEMM", spec=dict(_GLM, rows=2, in_kernel_norm=False), ), + Cell( + why="TP8 partial B1: 768 hidden-half outputs with fused hnorm", + spec=dict(_GLM, n=768, rows=1, in_kernel_norm=True), + ), + Cell( + why="TP8 full B2: 768 hidden-half outputs after weighted hnorm", + spec=dict(_GLM, n=768, rows=2, in_kernel_norm=False), + ), ) note = """ - The tactic is the shape's `splitk_default_tactic`; the in-kernel norm needs one row, H = - 6144 and a tactic whose CTA K slice fits its stages (`hnorm_b_supported`), else ValueError. + Plain GEMM uses the shape's `splitk_default_tactic`. The in-kernel norm needs one row, + H = 6144 and a tactic whose CTA K slice fits its stages (`hnorm_b_supported`); when the + TP8 default tile is too wide, `hidden_half_tactic` selects a 64-row tile with six stages. PDL follows `TRTLLM_ENABLE_PDL` (default on): before its grid dependency wait the kernel reads only the weight (TMA) and the norm weight and zeroes its own B stages; it reads `x` after the wait. @@ -91,8 +126,13 @@ def __call__( """ rows, k = x.shape sm_count, smem_capacity = device_limits(x.device.index) - tactic = splitk_default_tactic( - rows, weight.shape[0], k, sm_count=sm_count, smem_capacity=smem_capacity + tactic = hidden_half_tactic( + rows, + weight.shape[0], + k, + in_kernel_norm=norm_weight is not None, + sm_count=sm_count, + smem_capacity=smem_capacity, ) return run_splitk_hidden_half( x, weight, out, get_env_enable_pdl(), tactic, norm_weight=norm_weight, norm_eps=eps diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/fused_a_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/fused_a_gemm.py similarity index 79% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/fused_a_gemm.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/fused_a_gemm.py index 5f02dd788c34..b03a3bf38231 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/fused_a_gemm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/fused_a_gemm.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """GLM-5 fused-A GEMM (kv_a_proj_with_mqa, optionally with the DSA indexer's wk | weights_proj -rows) for up to 8 tokens, with an L2 prefetch of later weights, via the trtllm glm_fused_a_gemm +rows) for up to 32 tokens, with an L2 prefetch of later weights, via the trtllm glm_fused_a_gemm op.""" from __future__ import annotations @@ -16,13 +16,16 @@ true_fp32_matmul, ) +from .low_m_gemm import low_m_gemm + HIDDEN_SIZE = 6144 # kv_a_proj_with_mqa: q_a (2048) | compressed KV (512) | k_pe (64). OUT_FEATURES = 2624 # The same rows, then the indexer's wk (128) and weights_proj (32) rows. INDEXER_OUT_FEATURES = 2784 ROWS_PER_CTA = 16 -MAX_TOKENS = 8 +NATIVE_MAX_TOKENS = 8 +MAX_TOKENS = 32 # What a c1 decode step lists behind fused-A, at TP16 per-rank shapes. _INDEXER_WQ_B = ((4096, 2048), torch.bfloat16) # indexer wq_b, q_a norm folded in (not TP-split) @@ -41,7 +44,7 @@ class GlmFusedAGemm(OpWrapper): `[2560, 2624)` k_pe, `[2624, 2752)` indexer K, `[2752, 2784)` indexer weights; columns `[0, 2624)` are bitwise those of the 2624-row call. - One CTA per 16 weight rows (164 or 174 CTAs, one per SM) with all eight warps issuing the + Through eight token rows, one CTA per 16 weight rows (164 or 174 CTAs, one per SM) with all eight warps issuing the HMMAs. Before its grid-dependency wait the kernel stages the CTA's weight rows (tensor copies into shared memory, two thirds of them on into registers) and then prefetches the `l2_prefetch` tensors into L2; it reads `hidden_states` only after the wait and triggers its @@ -53,6 +56,14 @@ class GlmFusedAGemm(OpWrapper): # The c1 decode step runs the op 85 times: 78 verify layers at 8 rows plus 7 MTP draft steps. # 22 calls (21 indexer layers and draft step 0) take the 2784-row weight. CELLS: tuple[Cell, ...] = ( + *( + Cell( + why="TP4 verify fused-A through split-K", + spec=dict(tokens=m, out_features=n, l2_prefetch=()), + ) + for m in (12, 16, 17, 18) + for n in (2624, 2784) + ), Cell( why="indexer layers' verify step and MTP draft step 0: 2784 rows, 8 tokens (3D " "activation copies), the list indexer wq_b, v_b, o_proj", @@ -89,10 +100,11 @@ class GlmFusedAGemm(OpWrapper): PDL (`TRTLLM_ENABLE_PDL`): before its grid-dependency wait the kernel reads `weight`, so no kernel launched earlier in the step may write it (load-time weights only). Its dependents launch once it passes the wait and must wait on their own grid dependency before reading the - output or writing `hidden_states`. + output or writing `hidden_states`. Above eight rows the owned split-K GEMM + stages its own weights before the wait; the extra L2 hints are unused. Preconditions, enforced loudly by the op: CUDA BF16 2-D `hidden_states` and `weight` on one - device, contiguous and 16-byte aligned; hidden 6144; 2624 or 2784 rows; 1 to 8 tokens; at most + device, contiguous and 16-byte aligned; hidden 6144; 2624 or 2784 rows; 1 to 32 tokens; at most four `l2_prefetch` tensors, each contiguous, aligned and on the device; SM90 or later. Numerics: the products of a plain GEMM, summed in FP32 in the kernel's order (two @@ -107,7 +119,24 @@ def __call__( weight: torch.Tensor, l2_prefetch: list[torch.Tensor], ) -> torch.Tensor: - return torch.ops.trtllm.glm_fused_a_gemm(hidden_states, weight, l2_prefetch) + if hidden_states.shape[0] <= NATIVE_MAX_TOKENS: + return torch.ops.trtllm.glm_fused_a_gemm(hidden_states, weight, l2_prefetch) + if not 1 <= hidden_states.shape[0] <= MAX_TOKENS: + raise ValueError(f"fused-A supports at most {MAX_TOKENS} rows") + if ( + hidden_states.ndim != 2 + or hidden_states.shape[1] != HIDDEN_SIZE + or weight.shape + not in ((OUT_FEATURES, HIDDEN_SIZE), (INDEXER_OUT_FEATURES, HIDDEN_SIZE)) + or hidden_states.dtype != torch.bfloat16 + or weight.dtype != torch.bfloat16 + or not hidden_states.is_contiguous() + or not weight.is_contiguous() + ): + raise ValueError("unsupported fused-A weight or activation layout") + return low_m_gemm( + hidden_states, weight, hidden_states.new_empty(hidden_states.shape[0], weight.shape[0]) + ) def reference( self, @@ -142,7 +171,7 @@ def is_valid( assert weight.shape[0] in (OUT_FEATURES, INDEXER_OUT_FEATURES), ( f"weight must have {OUT_FEATURES} or {INDEXER_OUT_FEATURES} rows, got {weight.shape[0]}" ) - # One activation tile of 8 token rows per k stage. + # Native HMMA through eight rows, cluster split-K for wider batches. assert 1 <= hidden_states.shape[0] <= MAX_TOKENS, ( f"the kernel takes 1 to {MAX_TOKENS} tokens, got {hidden_states.shape[0]}" ) @@ -152,7 +181,9 @@ def is_valid( # instead (32-row tiles, half the CTAs, for the 2624-row weight). sms = torch.cuda.get_device_properties(hidden_states.device).multi_processor_count ctas = weight.shape[0] // ROWS_PER_CTA - assert sms >= ctas, f"{ctas} CTAs need as many SMs for one wave, the device has {sms} SMs" + assert hidden_states.shape[0] > NATIVE_MAX_TOKENS or sms >= ctas, ( + f"{ctas} CTAs need as many SMs for one wave, the device has {sms} SMs" + ) fused_a_gemm = GlmFusedAGemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/indexer_q_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/indexer_q_tail.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/indexer_q_tail.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/indexer_q_tail.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/low_m_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/low_m_gemm.py similarity index 77% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/low_m_gemm.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/low_m_gemm.py index a232199870d6..f8e4235dc528 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/low_m_gemm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/low_m_gemm.py @@ -14,7 +14,13 @@ ) from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl -from ._low_m_splitk import device_limits, operands_overlap, run_splitk_dense, splitk_default_tactic +from ._low_m_splitk import ( + device_limits, + one_wave_tactic, + operands_overlap, + run_splitk_dense, + splitk_default_tactic, +) class GlmLowMGemm(OpWrapper): @@ -62,6 +68,23 @@ class GlmLowMGemm(OpWrapper): why="eh_proj, MTP draft step 0: 8 x 768 (row-parallel slice) -> 6144, split-K 2", spec=dict(m=8, k=768, n=6144), ), + *( + Cell(why=f"TP4 projection M={m} N={n} K={k}", spec=dict(m=m, n=n, k=k)) + for m in (1, 3, 6, 12, 18) + for n, k in ((1024, 6144), (6144, 512), (38720, 6144), (6144, 3072)) + ), + *( + Cell(why=f"TP8 projection M={m} N={n} K={k}", spec=dict(m=m, n=n, k=k)) + for m in (1, 2, 4, 8) + for n, k in ( + (3072, 6144), + (6144, 1536), + (512, 6144), + (6144, 256), + (19360, 6144), + (6144, 2048), + ) + ), ) note = """ @@ -77,15 +100,23 @@ class GlmLowMGemm(OpWrapper): kernel; make it before CUDA-graph capture. """ - def __call__(self, x: torch.Tensor, weight: torch.Tensor, out: torch.Tensor) -> torch.Tensor: + def __call__( + self, x: torch.Tensor, weight: torch.Tensor, out: torch.Tensor, *, one_wave: bool = False + ) -> torch.Tensor: """x [M, K] BF16 row-major; weight [N, K] BF16 (the stored weight); out [M, N] BF16.""" m, k = x.shape n = weight.shape[0] sm_count, smem_capacity = device_limits(x.device.index) tactic = splitk_default_tactic(m, n, k, sm_count=sm_count, smem_capacity=smem_capacity) + if one_wave: + tactic = one_wave_tactic( + tactic, m, n, k, sm_count=sm_count, smem_capacity=smem_capacity + ) return run_splitk_dense(x, weight.t(), None, out, get_env_enable_pdl(), tactic) - def reference(self, x: torch.Tensor, weight: torch.Tensor, out: torch.Tensor) -> torch.Tensor: + def reference( + self, x: torch.Tensor, weight: torch.Tensor, out: torch.Tensor, *, one_wave: bool = False + ) -> torch.Tensor: """The product the op writes into `out`: FP32-accumulated (TF32 off), rounded once. `out` is named to mirror `__call__` and not read. @@ -93,7 +124,9 @@ def reference(self, x: torch.Tensor, weight: torch.Tensor, out: torch.Tensor) -> with true_fp32_matmul(): return (x.float() @ weight.float().t()).to(x.dtype) - def is_valid(self, x: torch.Tensor, weight: torch.Tensor, out: torch.Tensor) -> None: + def is_valid( + self, x: torch.Tensor, weight: torch.Tensor, out: torch.Tensor, *, one_wave: bool = False + ) -> None: # The owner CTA of a tile stores its rows while other clusters still stream x and the # weight; an `out` that shares memory with either is overwritten under their reads and # the result is silently wrong. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/o_proj_publish.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/o_proj_publish.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/o_proj_publish.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/o_proj_publish.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/q_b_fp8_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/q_b_fp8_rope.py similarity index 96% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/q_b_fp8_rope.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/q_b_fp8_rope.py index 44637c3ea359..cef020ee9ba7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/q_b_fp8_rope.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/q_b_fp8_rope.py @@ -72,6 +72,14 @@ class GlmQBFp8Rope(OpWrapper): why="two requests of 8 tokens: per-sequence positions, 16-row tiles (regular plan)", spec=dict(_GLM, m=16, seq_len=8, plan="regular"), ), + *( + Cell( + why=f"TP4 absorbed q_b: {m} rows, {plan}", + spec=dict(_GLM, heads=16, m=m, seq_len=6 if m >= 6 else 1, plan=plan), + ) + for m in (1, 3, 6, 12, 18) + for plan in ("regular", "beside_topk", "early_gather") + ), ) note = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/router_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/router_gemm.py similarity index 84% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/router_gemm.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/router_gemm.py index 990fe197a7ea..bdf1943b9c8f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/gemm/router_gemm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/gemm/router_gemm.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GLM-5 MoE router logits (256 experts at hidden 6144, FP32 from BF16) for up to 16 tokens, via +"""GLM-5 MoE router logits (256 experts at hidden 6144, FP32 from BF16) for up to 32 tokens, via the trtllm glm_router_gemm op.""" from __future__ import annotations @@ -12,13 +12,14 @@ HIDDEN_SIZE = 6144 NUM_EXPERTS = 256 -MAX_TOKENS = 16 +MAX_TOKENS = 32 class GlmRouterGemm(OpWrapper): """`hidden_states @ weight.T` in FP32 from BF16 operands: the router logits `[tokens, 256]`. - `weight` is the gate weight as stored, `[256, 6144]` BF16 row-major. One CTA per expert. Before + `weight` is the gate weight as stored, `[256, 6144]` BF16 row-major. Through16 tokens there + is one CTA per expert;17-32 use eight-CTA clusters. Before its grid-dependency wait each CTA prefetches its weight row into L2; right after the wait, before any logit is written, it lets its dependents launch. """ @@ -36,7 +37,11 @@ class GlmRouterGemm(OpWrapper): why="8 rows: a verify-width step routed without the row select", spec=dict(tokens=8), ), - Cell(why="16 tokens: the op's ceiling, its largest instance", spec=dict(tokens=16)), + Cell(why="16 tokens: the largest GEMV instance", spec=dict(tokens=16)), + *( + Cell(why=f"TP4 tensorcore router: {tokens} rows", spec=dict(tokens=tokens)) + for tokens in (17, 18, 24, 32) + ), ) note = """ @@ -48,11 +53,13 @@ class GlmRouterGemm(OpWrapper): before its own wait, which holds because FC1 releases it only after FC1's wait. Preconditions, enforced loudly by the op: CUDA BF16 2-D `hidden_states` and `weight` on one - device, contiguous and 16-byte aligned; hidden 6144; 256 experts; 1 to 16 tokens. + device, contiguous and 16-byte aligned; hidden 6144; 256 experts; 1 to 32 tokens. Numerics: per thread an FP32 FMA chain over its 48 products, then a warp butterfly and the four warp sums in order; deterministic, and bitwise the logits of the shared - `dsv3_router_gemm_op`, whose kernel sums in the same order. + `dsv3_router_gemm_op` through 16 tokens. Above 16, an eight-CTA cluster uses + BF16 tensor cores and a fixed warp-then-cluster reduction order; FP32 logits + approximate the same product, but are not bitwise the GEMV reduction. """ def __call__(self, hidden_states: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: @@ -68,7 +75,7 @@ def reference(self, hidden_states: torch.Tensor, weight: torch.Tensor) -> torch. def is_valid(self, hidden_states: torch.Tensor, weight: torch.Tensor) -> None: # The shape gates of the kernel, which a target checks to route anything else to the - # shared dsv3_router_gemm_op. The op raises on them as well. + # dense router fallback. The op raises on them as well. assert hidden_states.dim() == 2 and hidden_states.shape[1] == HIDDEN_SIZE, ( f"hidden_states must be [tokens, {HIDDEN_SIZE}]" ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/moe/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/moe/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/lowm_bf16_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/moe/lowm_bf16_moe.py similarity index 65% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/lowm_bf16_moe.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/moe/lowm_bf16_moe.py index f2b83f9f4ede..11b4bb7ffb87 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/lowm_bf16_moe.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/moe/lowm_bf16_moe.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""One token through the MTP layer's BF16 MoE, routing included, via the glm_lowm_bf16_moe op.""" +"""Draft tokens through the MTP layer's BF16 MoE, routing included, via the glm_lowm_bf16_moe op.""" from __future__ import annotations @@ -16,7 +16,7 @@ class GlmLowmBf16Moe(OpWrapper): - """The routed half of the MTP layer's MoE for one token, from the router logits to the weighted expert sum. + """The routed half of the MTP layer's MoE for one to six tokens, from the router logits to the weighted expert sum. Routing is the n_group == 1 case of `noaux_tc_op`: experts by `sigmoid(logits) + bias` (sigmoid in the tanh form), the top 8 in descending order with the lower expert id first on ties, weights from the unbiased scores @@ -25,7 +25,7 @@ class GlmLowmBf16Moe(OpWrapper): `bf16(sum_k bf16(w_k * (W2[e_k] act_k)))`, the sum in FP32 in the routing order. `w3_w1` is `[256, 2 * I, H]` with the up rows (W3) first and the gate rows (W1) second, `w2` is `[256, H, I]`, - all experts local (tensor-parallel slices, no expert parallelism). Returns `[1, H]`; the shared experts are not + all experts local (tensor-parallel slices, no expert parallelism). Returns `[M, H]`; the shared experts are not part of the op. """ @@ -47,11 +47,28 @@ class GlmLowmBf16Moe(OpWrapper): why="selected scores summing below 2^-13: the weights take the FP64 expression", spec=dict(hidden=6144, inter=128, routing="tiny"), ), + *( + Cell( + why=f"TP4 draft batch {tokens}: intermediate512, {routing} routing", + spec=dict(tokens=tokens, hidden=6144, inter=512, routing=routing), + ) + for tokens in (1, 2, 3) + for routing in ("random", "ties", "tiny") + ), + *( + Cell( + why=f"TP8 draft batch {tokens}: intermediate256, {routing} routing", + spec=dict(tokens=tokens, hidden=6144, inter=256, routing=routing), + ) + for tokens in (1, 2) + for routing in ("random", "ties", "tiny") + ), ) note = """ The kernels serve hidden sizes that are multiples of 2048 up to 8192 and expert slices that are multiples of - 64 up to 512, exactly 256 experts and top 8; the op refuses anything else, and more than one token, loudly. A + 64 up to 512, exactly 256 experts and top 8. Multi-token calls serve 2--6 rows with + hidden size at most 6144; the target uses the wide512 slice only for draft batches2/3. A BF16 correction bias is accepted as well; the checkpoint's is FP32, so no cell drives it. The order of the two halves of `w3_w1` is not visible in the shapes: a gate-first tensor gives silently wrong @@ -91,25 +108,28 @@ def reference( `noaux_tc_op` forms them, the expert sum in FP32 in the routing order. Products and sums are FP32 (no TF32). """ - inter = w2.shape[2] - scores = 0.5 * torch.tanh(0.5 * router_logits.float()[0]) + 0.5 - choice = scores + e_score_correction_bias.float() - ids = torch.sort(choice, descending=True, stable=True).indices[:top_k] - selected = scores[ids] - total = selected.sum().double() - weights = (selected.double() * routed_scaling_factor / (total + 1e-20)).float() - xf = x.float()[0] - out = torch.zeros(w2.shape[1], dtype=torch.float32, device=x.device) - with true_fp32_matmul(): - for k in range(top_k): - expert = int(ids[k]) - fc1 = w3_w1[expert].float() @ xf - up = fc1[:inter].to(torch.bfloat16).float() - gate = fc1[inter:].to(torch.bfloat16).float() - act = (gate * (1.0 / (1.0 + torch.exp(-gate))) * up).to(torch.bfloat16).float() - term = (weights[k] * (w2[expert].float() @ act)).to(torch.bfloat16).float() - out = out + term - return out.to(torch.bfloat16).unsqueeze(0) + rows = [] + for row in range(x.shape[0]): + inter = w2.shape[2] + scores = 0.5 * torch.tanh(0.5 * router_logits.float()[row]) + 0.5 + choice = scores + e_score_correction_bias.float() + ids = torch.sort(choice, descending=True, stable=True).indices[:top_k] + selected = scores[ids] + total = selected.sum().double() + weights = (selected.double() * routed_scaling_factor / (total + 1e-20)).float() + xf = x.float()[row] + out = torch.zeros(w2.shape[1], dtype=torch.float32, device=x.device) + with true_fp32_matmul(): + for k in range(top_k): + expert = int(ids[k]) + fc1 = w3_w1[expert].float() @ xf + up = fc1[:inter].to(torch.bfloat16).float() + gate = fc1[inter:].to(torch.bfloat16).float() + act = (gate * (1.0 / (1.0 + torch.exp(-gate))) * up).to(torch.bfloat16).float() + term = (weights[k] * (w2[expert].float() @ act)).to(torch.bfloat16).float() + out = out + term + rows.append(out.to(torch.bfloat16)) + return torch.stack(rows) def is_valid( self, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/routed_experts.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/moe/routed_experts.py similarity index 89% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/routed_experts.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/moe/routed_experts.py index c747ff6e9e8e..9f0b770e0839 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/moe/routed_experts.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/moe/routed_experts.py @@ -6,7 +6,7 @@ `trtllm::glm_routed_experts`: the TRTLLM-Gen MoE kernels with GLM's own router and FC2 kernels beside them. -On the decode step (up to 8 rows) the call runs the router GEMV inside the +On the decode step (up to 24 rows) the call runs the router GEMV inside the routing kernel, lets the NVFP4 input still be in flight on another stream until FC1 starts, and runs FC2 on the K = 128 kernel; it stops before the finalize step, which the post-MoE all-reduce hop performs.""" @@ -52,7 +52,7 @@ ACT_TYPE_SWIGLU = 0 #: The largest row count the fused router + routing kernel takes. -FUSED_ROUTER_MAX_TOKENS = 8 +FUSED_ROUTER_MAX_TOKENS = 24 #: One expert stack's `(codes [E, N, K] uint8, block scales [E, N, K/16] #: float8_e4m3fn)`, unshuffled and unswizzled. @@ -97,7 +97,7 @@ def __init__( act_type: int, tune_max_num_tokens: int, use_dp: bool, - prefer_gemm2_k128: bool, + prefer_specialized_gemm2: bool, ): super().__init__( num_experts, @@ -114,12 +114,12 @@ def __init__( tune_max_num_tokens=tune_max_num_tokens, use_dp=use_dp, ) - self.prefer_gemm2_k128 = prefer_gemm2_k128 + self.prefer_specialized_gemm2 = prefer_specialized_gemm2 def unique_id(self): - # FC2 on the K = 128 kernel times the tile-8 tactics differently. + # Owned FC2 kernels time the tile-8 tactics differently from the generic cubin. key = super().unique_id() - return key + ("gemm2_k128",) if self.prefer_gemm2_k128 else key + return key + ("specialized_gemm2",) if self.prefer_specialized_gemm2 else key def get_runner(self): key = (self.act_type,) @@ -175,7 +175,7 @@ def forward( router_hidden_states, router_weight, router_workspace, - self.prefer_gemm2_k128, + self.prefer_specialized_gemm2, ) @@ -217,7 +217,7 @@ def _glm_routed_experts( router_hidden_states: Optional[torch.Tensor], router_weight: Optional[torch.Tensor], router_workspace: Optional[torch.Tensor], - prefer_gemm2_k128: bool, + prefer_specialized_gemm2: bool, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """The op behind `GlmRoutedExperts.__call__` (see there). The autotuner profiles it routed from logits, as it profiles the TRTLLM-Gen FP4 MoE op.""" @@ -236,7 +236,7 @@ def _glm_routed_experts( act_type, tune_max_num_tokens, use_dp, - prefer_gemm2_k128, + prefer_specialized_gemm2, ) routing_logits_for_tuner, topk_weights_for_tuner, topk_ids_for_tuner, tuning_config = ( prepare_dummy_topk_and_hook( @@ -340,7 +340,7 @@ def _( router_hidden_states, router_weight, router_workspace, - prefer_gemm2_k128, + prefer_specialized_gemm2, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if do_finalize: finalized = ( @@ -372,22 +372,22 @@ def _dequantize_activation(hidden_states: torch.Tensor, hidden_states_scale: tor class GlmRoutedExperts(OpWrapper): - """One routed-expert MoE layer of GLM-5.2 at MoE TP16, over NVFP4 input. + """One routed-expert MoE layer of GLM-5.2 at MoE TP4/TP16, over NVFP4 input. Routing happens inside the op, from one of two sources: * `routing_logits` `[T, 256]` -- the router GEMM's fp32 logits; * `router_hidden_states` `[T, 6144]` bf16 with `router_weight` `[256, 6144]` bf16 and `router_workspace` (`make_fused_router_workspace`): - the routing kernel runs the router GEMV itself, bit for bit the logits of - `dsv3_router_gemm_op`, at 1-8 rows. `routing_logits` then only gives the + the routing kernel computes the router product itself: GEMV through16 rows, + tensor-core GEMM at17-24. `routing_logits` then only gives the shape and is not read. `hidden_states_ready_event` is the raw `cudaEvent_t` of a `torch.cuda.Event` recorded on the stream that writes `hidden_states` / `hidden_states_scale` (`event.cuda_event`, after its first record): routing launches without - waiting for it, FC1 waits. `prefer_gemm2_k128` runs FC2 on the persistent - K = 128 kernel where it covers the call -- intermediate 128, tile N 8 and + waiting for it, FC1 waits. `prefer_specialized_gemm2` runs FC2 on the persistent + K = 128 or K = 256 kernel where it covers the call -- matching intermediate, tile N 8 and enough SMs, which an 8-row call meets on Rubin's SM count -- with the Gemm2 cubin's rows, bitwise. @@ -415,17 +415,73 @@ class GlmRoutedExperts(OpWrapper): Cell( why="the c1 verify step: 8 rows routed by the fused router GEMV, the input joined " "right before FC1, FC2 at K = 128, unfinalized for the post-MoE hop", - spec=dict(tokens=8, fused_router=True, late_join=True, gemm2_k128=True, **_GLM_TP16), + spec=dict( + tokens=8, fused_router=True, late_join=True, specialized_gemm2=True, **_GLM_TP16 + ), ), Cell( why="one row, a step without MTP: the fused router's single-token instance and " "FC2 at K = 128 over 8 token tiles", - spec=dict(tokens=1, fused_router=True, late_join=True, gemm2_k128=True, **_GLM_TP16), + spec=dict( + tokens=1, fused_router=True, late_join=True, specialized_gemm2=True, **_GLM_TP16 + ), ), Cell( why="32 rows (4 requests x 8), past the fused router: routed from the router " "GEMM's logits, FC2 back on the Gemm2 cubin", - spec=dict(tokens=32, fused_router=False, late_join=False, gemm2_k128=True, **_GLM_TP16), + spec=dict( + tokens=32, fused_router=False, late_join=False, specialized_gemm2=True, **_GLM_TP16 + ), + ), + Cell( + why="TP4 verify 6 rows with fused routing and intermediate512", + spec=dict( + tokens=6, + fused_router=True, + late_join=True, + specialized_gemm2=False, + **{**_GLM_TP16, "intermediate": 512}, + ), + ), + Cell( + why="TP4 verify 12 rows with fused routing and intermediate512", + spec=dict( + tokens=12, + fused_router=True, + late_join=True, + specialized_gemm2=False, + **{**_GLM_TP16, "intermediate": 512}, + ), + ), + Cell( + why="TP4 verify 18 rows with fused routing and intermediate512", + spec=dict( + tokens=18, + fused_router=True, + late_join=True, + specialized_gemm2=False, + **{**_GLM_TP16, "intermediate": 512}, + ), + ), + Cell( + why="TP8 verify 4 rows with fused routing and FC2 K256", + spec=dict( + tokens=4, + fused_router=True, + late_join=True, + specialized_gemm2=True, + **{**_GLM_TP16, "intermediate": 256}, + ), + ), + Cell( + why="TP8 verify 8 rows with fused routing and FC2 K256", + spec=dict( + tokens=8, + fused_router=True, + late_join=True, + specialized_gemm2=True, + **{**_GLM_TP16, "intermediate": 256}, + ), ), ) @@ -496,7 +552,7 @@ def __call__( router_hidden_states: Optional[torch.Tensor], router_weight: Optional[torch.Tensor], router_workspace: Optional[torch.Tensor], - prefer_gemm2_k128: bool, + prefer_specialized_gemm2: bool, ) -> list[torch.Tensor]: result = torch.ops.trtllm.glm_routed_experts( routing_logits, @@ -535,7 +591,7 @@ def __call__( router_hidden_states, router_weight, router_workspace, - prefer_gemm2_k128, + prefer_specialized_gemm2, ) return [result[0]] if do_finalize else list(result) @@ -657,7 +713,7 @@ def is_valid( router_hidden_states: Optional[torch.Tensor], router_weight: Optional[torch.Tensor], router_workspace: Optional[torch.Tensor], - prefer_gemm2_k128: bool, + prefer_specialized_gemm2: bool, ) -> None: """The ways this op takes an input and answers it wrongly. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/norm/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/norm/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/norm/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/quantization/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/quantization/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/quantization/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/quantization/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/quantization/fp4_quantize.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/quantization/fp4_quantize.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/quantization/fp4_quantize.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/quantization/fp4_quantize.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_exchange.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/_draft_exchange.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_exchange.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/_draft_exchange.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_glue_rmsnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/_draft_glue_rmsnorm.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_glue_rmsnorm.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/_draft_glue_rmsnorm.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_step_state.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/_draft_step_state.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_step_state.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/_draft_step_state.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_token_table.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/_draft_token_table.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_draft_token_table.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/_draft_token_table.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_gather_last_accepted_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/_gather_last_accepted_rows.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/_gather_last_accepted_rows.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/_gather_last_accepted_rows.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_ehproj_table.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/draft_ehproj_table.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_ehproj_table.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/draft_ehproj_table.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_glue_rmsnorm_cat.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/draft_glue_rmsnorm_cat.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_glue_rmsnorm_cat.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/draft_glue_rmsnorm_cat.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_sample_glue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/draft_sample_glue.py similarity index 92% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_sample_glue.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/draft_sample_glue.py index 741c387dedcf..871fd713d967 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/draft_sample_glue.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/draft_sample_glue.py @@ -38,7 +38,7 @@ def _step(**overrides) -> dict: class GlmDraftSampleGlue(OpWrapper): - """The greedy draft token of every row, and optionally the next step's inputs and state, in one CTA. + """Greedy draft tokens and next-step inputs/state, in one CTA (eight CTAs at TP4). Collective over the TP group (`group_size` ranks, the workspace's group): each rank takes the argmax of its logits shard `[rows, vocab_shard]` (rank r holds tokens `[r * vocab_shard, (r + 1) * vocab_shard)`), the ranks @@ -89,6 +89,34 @@ class GlmDraftSampleGlue(OpWrapper): why="the most rows one launch serves (8 requests), last step", spec=_step(rows=8, next_new=True), ), + *( + Cell( + why=f"TP4 K5 batch{rows} step0: wide argmax, gathered positions and rewind", + spec=_step( + rows=rows, + draft_len=5, + step=0, + vocab_shard=38720, + glue=True, + positions="gather", + fill=True, + kv_update=2, + ), + ) + for rows in (1, 2, 3) + ), + Cell( + why="TP4 K5 batch3 middle step: wide argmax and eh_proj input", + spec=_step( + rows=3, draft_len=5, step=2, vocab_shard=38720, glue=True, positions="plain" + ), + ), + Cell( + why="TP4 K5 batch3 last step: tied argmax and next-new-token rows", + spec=_step( + rows=3, draft_len=5, step=4, vocab_shard=38720, logits="tied", next_new=True + ), + ), ) note = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/gather_last_accepted_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/gather_last_accepted_rows.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/gather_last_accepted_rows.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/gather_last_accepted_rows.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/verify_head_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/verify_head_accept.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/ops/spec/verify_head_accept.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/verify_head_accept.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/verify_head_sample.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/verify_head_sample.py new file mode 100644 index 000000000000..d2ac516156a2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/ops/spec/verify_head_sample.py @@ -0,0 +1,178 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tokens-only FULL sampling for gathered Rubin GLM verify logits. + +Logits may be finite or -inf (masked), with a finite scaled maximum in each row. Temperature, +filters and RNG arrays are completed step inputs before the PDL logits producer. +The state is stream ordered and cannot serve concurrent CUDA streams. Native +scratch and counters are real mutated tensors; no indirect allocation is hidden. +""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +VOCAB_SIZE = 154880 +MAX_ROWS = 24 + + +def workspace_bytes(rows: int, vocab: int = VOCAB_SIZE) -> int: + return ((4 + 32) * ((vocab + 511) // 512) + 4096) * rows * 4 + + +def _supported(device, max_rows, logits, temperatures, top_ks, top_ps, min_ps, seed, offset): + if ( + not isinstance(logits, torch.Tensor) + or logits.device != device + or logits.dtype != torch.float32 + or logits.ndim != 2 + or not 1 <= logits.shape[0] <= max_rows + or logits.shape[1] != VOCAB_SIZE + or not logits.is_contiguous() + or logits.data_ptr() % 16 + ): + return False + rows = logits.shape[0] + for value, dtype, count in ( + (temperatures, torch.float32, rows), + (top_ks, torch.int32, rows), + (top_ps, torch.float32, rows), + (min_ps, torch.float32, rows), + (seed, torch.int64, 1), + (offset, torch.int64, 1), + ): + if ( + not isinstance(value, torch.Tensor) + or value.device != device + or value.dtype != dtype + or value.ndim != 1 + or not value.is_contiguous() + or value.numel() < count + ): + return False + return seed.numel() == offset.numel() and (seed.numel() == 1 or seed.numel() >= rows) + + +class GlmVerifyHeadSample(OpWrapper): + """Sample FULL weights after min-p, tie-preserving top-k and inclusive top-p. + + The reference returns the categorical distribution, not an unrelated RNG's + sampled token. Tests compare support/distribution and this op's own determinism. + Different reduction order/random consumption can change IDs versus the shared + sampler, even with identical per-row Philox initialization. + """ + + ARCHS = frozenset({Arch.SM_107}) + CELLS = tuple( + Cell( + why=f"GLM verify rows={rows}, FULL filters and per-token RNG", + spec=dict(rows=rows, vocab=VOCAB_SIZE), + ) + for rows in (1, 6, 12, 18, 24) + ) + note = "Masked -inf logits allowed with finite row maxima; no cross-algorithm ID equality." + + def __call__( + self, logits, temperatures, top_ks, top_ps, min_ps, seed, offset, row_counters, workspace + ): + return torch.ops.trtllm.glm_verify_head_sample_tokens( + logits, temperatures, top_ks, top_ps, min_ps, seed, offset, row_counters, workspace + ) + + def reference( + self, logits, temperatures, top_ks, top_ps, min_ps, seed, offset, row_counters, workspace + ): + """Independent categorical probabilities; RNG and scratch do not affect the law. + + Weight scaling is FP32; sorted float64 accumulation makes boundary tests + use their specified mathematical mass, avoiding a duplicate CUDA reduction. + """ + inv_temp = torch.where( + temperatures[: len(logits)] > 0, 1.0 / temperatures[: len(logits)], 1.0 / 1.0e-6 + ) + scaled = logits * inv_temp[:, None] + weights = (scaled - scaled.amax(dim=-1, keepdim=True)).exp() + result = [] + for row, weight in enumerate(weights): + floor = min_ps[row].clamp(0, 1) + k = int(top_ks[row]) + if 0 < k < weight.numel(): + floor = torch.maximum(floor, weight.topk(k).values[-1]) + weight = torch.where(weight >= floor, weight, 0.0) + if float(top_ps[row]) < 1: + ordered = weight.sort(descending=True).values + cumulative = ordered.double().cumsum(0) + # First index of each tie gives the mass of STRICTLY greater weights. + first = torch.searchsorted(-ordered, -weight, right=False) + prefix = torch.cat((cumulative.new_zeros(1), cumulative)) + above = prefix[first] + weight = torch.where(above <= top_ps[row].double() * cumulative[-1], weight, 0.0) + result.append(weight / weight.sum()) + return torch.stack(result) + + def is_valid( + self, logits, temperatures, top_ks, top_ps, min_ps, seed, offset, row_counters, workspace + ): + assert logits.is_cuda and _supported( + logits.device, MAX_ROWS, logits, temperatures, top_ks, top_ps, min_ps, seed, offset + ) + assert not torch.isnan(logits).any() and not torch.isposinf(logits).any() + assert torch.isfinite(temperatures).all() + inv_temp = torch.where( + temperatures[: len(logits)] > 0, 1.0 / temperatures[: len(logits)], 1.0 / 1.0e-6 + ) + assert torch.isfinite((logits * inv_temp[:, None]).amax(dim=1)).all() + assert torch.isfinite(top_ps).all() and ((top_ps >= 0) & (top_ps <= 1)).all() + assert torch.isfinite(min_ps).all() + assert row_counters.dtype == torch.int32 and row_counters.device == logits.device + assert row_counters.ndim == 1 and row_counters.numel() >= len(logits) + assert row_counters.is_contiguous() and not row_counters.any() + assert workspace.dtype == torch.uint8 and workspace.device == logits.device + assert workspace.ndim == 1 and workspace.is_contiguous() + assert workspace.numel() >= workspace_bytes(len(logits)) and workspace.data_ptr() % 16 == 0 + + +glm_verify_head_sample = GlmVerifyHeadSample() + + +class VerifyHeadSampleState: + """Persistent scratch allocated before profiling/capture; one ordered execution stream.""" + + def __init__(self, device, max_rows: int = MAX_ROWS, vocab: int = VOCAB_SIZE): + device = torch.device(device) + if device.type != "cuda" or not 1 <= max_rows <= MAX_ROWS or vocab != VOCAB_SIZE: + raise ValueError("GLM verify sampling requires CUDA, rows1..24 and vocabulary154880") + if device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + if torch.cuda.get_device_capability(device) != (10, 7): + raise ValueError("GLM verify sampling requires SM107") + self.device = device + self.max_rows = max_rows + self.row_counters = torch.zeros(max_rows, dtype=torch.int32, device=device) + self.workspace = torch.empty(workspace_bytes(max_rows), dtype=torch.uint8, device=device) + + def supported(self, logits, temperatures, top_ks, top_ps, min_ps, seed, offset) -> bool: + # Shape/layout-only eligibility is stable under CUDA graph parameter updates. + if torch.compiler.is_compiling(): + return False + return _supported( + self.device, self.max_rows, logits, temperatures, top_ks, top_ps, min_ps, seed, offset + ) + + def sample(self, logits, temperatures, top_ks, top_ps, min_ps, seed, offset): + if not self.supported(logits, temperatures, top_ks, top_ps, min_ps, seed, offset): + raise ValueError("unsupported GLM verify sampling inputs; select shared sampler") + return glm_verify_head_sample( + logits, + temperatures, + top_ks, + top_ps, + min_ps, + seed, + offset, + self.row_counters, + self.workspace, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/weights.py similarity index 98% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/weights.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/weights.py index 098b69f41c4d..8ff0939d6dfc 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/weights.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/weights.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GLM-5.2-NVFP4 checkpoint tensors mapped onto the tp16 target's parameters. +"""GLM-5.2-NVFP4 checkpoint tensors mapped onto the Rubin target's parameters. The target owns its GLM-5.2 modules and loader. The checkpoint loads into the same shapes and layouts the context server uses. The MTP draft steps derive diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py index bfe5cf42faf4..0f69d37c883c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/routing.py @@ -47,14 +47,14 @@ def route( pretrained_config.q_lora_rank, ) index_n_heads = getattr(pretrained_config, "index_n_heads", None) - # tp16: attention heads and the expert intermediate dimension sharded over - # 16 ranks, no attention DP, no expert parallelism. A weight-layout + # pure TP: attention heads and the expert intermediate dimension sharded over + # 4, 8, or 16 ranks, no attention DP, no expert parallelism. A weight-layout # question, which is why it selects a target rather than a runtime branch. - tp16 = ( - mapping.world_size == 16 - and mapping.tp_size == 16 + pure_tp = ( + mapping.tp_size in (4, 8, 16) + and mapping.world_size == mapping.tp_size and mapping.moe_ep_size == 1 - and mapping.moe_tp_size == 16 + and mapping.moe_tp_size == mapping.tp_size and not mapping.enable_attention_dp and mapping.pp_size == 1 and mapping.cp_size == 1 @@ -69,11 +69,11 @@ def route( and getattr(spec_config, "eagle_choices", None) is None and spec_config.max_total_draft_tokens == spec_config.max_draft_len ) - # GLM-5.2-NVFP4 on 16 Rubin GPUs (sm107), tp16, MTP-Eagle or no speculation. + # GLM-5.2-NVFP4 on Rubin (sm107), TP4/TP8/TP16, MTP-Eagle or no speculation. if ( compute_capability == (10, 7) and shape == _SHAPE - and tp16 + and pure_tp and mtp_eagle_or_none and kv_cache_quant_algo == QuantAlgo.FP8 ): @@ -81,10 +81,10 @@ def route( # `from ..modeling import ` lines. # fmt: off # isort: off - from .glm_5_2_nvfp4__sm_107__tp16.modeling import ModelingV2Glm52Nvfp4Sm107Tp16 + from .glm_5_2_nvfp4__sm_107__low_latency.modeling import ModelingV2Glm52Nvfp4Sm107LowLatency # isort: on # fmt: on - return ModelingV2Glm52Nvfp4Sm107Tp16 + return ModelingV2Glm52Nvfp4Sm107LowLatency # tep4: the attention heads and the dense MLPs sharded four ways, the experts # sharded four ways. Attention DP would replicate attention, which is a @@ -142,7 +142,7 @@ def route( f"moe_tp={mapping.moe_tp_size} attention_dp={mapping.enable_attention_dp} " f"pp={mapping.pp_size} cp={mapping.cp_size}, speculative decoding {spec}; " f"KV cache {kv_cache_quant_algo or 'bf16'}; " - "the glm_moe_dsa targets are GLM-5.2-NVFP4 on sm107, tp16 (16 ranks, " + "the glm_moe_dsa targets are GLM-5.2-NVFP4 on sm107, TP4/TP8/TP16 (4/8/16 ranks, " "tensor-parallel attention and experts, no attention DP), with one-model " "MTP-Eagle as a linear chain or no speculative decoding; and GLM-5.2-NVFP4 " "on sm103, tep4 (4 ranks, tensor-parallel attention, expert-parallel MoE), " diff --git a/tests/integration/defs/.test_durations b/tests/integration/defs/.test_durations index f2f41d13c7da..71ec0a458c04 100644 --- a/tests/integration/defs/.test_durations +++ b/tests/integration/defs/.test_durations @@ -1306,8 +1306,8 @@ "test_unittests.py::test_unittests_v2[unittest/_torch/modeling/test_qwen4_exp_support.py]": 20.589279141104296, "unittest/_torch/modeling_v2/comm": 22.324, "test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/comm]": 22.324, - "unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_exchange.py": 30.08225925925926, - "test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_exchange.py]": 30.08225925925926, + "unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_draft_exchange.py": 30.08225925925926, + "test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_draft_exchange.py]": 30.08225925925926, "unittest/_torch/modeling_v2/models/glm_moe_dsa/test_routing.py": 25.9784, "test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/models/glm_moe_dsa/test_routing.py]": 25.9784, "unittest/_torch/modeling_v2/test_modeling_v2_claims.py": 24.492159420289855, diff --git a/tests/integration/defs/accuracy/modeling_v2/test_glm_moe_dsa.py b/tests/integration/defs/accuracy/modeling_v2/test_glm_moe_dsa.py index ac3b82ccca8d..6d5b8ff860bc 100644 --- a/tests/integration/defs/accuracy/modeling_v2/test_glm_moe_dsa.py +++ b/tests/integration/defs/accuracy/modeling_v2/test_glm_moe_dsa.py @@ -12,30 +12,21 @@ # 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. -"""Whole-model gates for the glm_moe_dsa modeling_v2 targets. +"""Whole-model accuracy and MTP coverage for the GLM SM103 and SM107 targets. -One file per model family, beside the other accuracy suites rather than inside -test_glm52.py: these gate a parallel implementation, and reading them next to -the built-in model's tests would invite treating one as a variant of the other. - -Every test here passes ``modeling_v2="require"`` to ``LLM``. Under ``"auto"`` a -configuration that missed a target's criteria -- or fell outside its bounds -- -would quietly fall back to the built-in implementation, pass, and report the -built-in's numbers as the target's. - -``TestModelingV2Glm52Nvfp4Sm107Tp16``: aggregated TP16/c1 accuracy and MTP -coverage for the SM107 target. Run with 16 MPI ranks on Rubin; the pytest -temporary directory must be visible to every rank. A short GPU trace proves -actual optimized decode and draft execution after engine warmup, including -CUDA graph replay. The full GSM8K evaluation then runs without profiling. - -``TestModelingV2Glm52Nvfp4Sm103Tep4``: identity and MTP1 GSM8K, paired in one -session, and MTP1 acceptance for the SM103 tep4 target. +Run each Rubin case with its 16, 8 or 4 MPI ranks; the pytest temporary directory must be visible +to every rank. A short GPU trace proves actual optimized decode and draft +execution after engine warmup, including CUDA graph replay. The full GSM8K +evaluation then runs without profiling. """ +import hashlib import json import os +import re import time +from collections import Counter +from dataclasses import asdict, dataclass from pathlib import Path import pytest @@ -63,6 +54,144 @@ ) from ._env import scoped_env_overrides +# The targets assert their own SM at construction: certification is per GPU +# architecture, and a receipt from another one says nothing here. +skip_not_sm103 = pytest.mark.skipif( + get_sm_version() != 103, reason="modeling_v2 targets in this batch are certified on sm_103 only" +) + + +class TestModelingV2Glm52Nvfp4Sm103Tep4(LlmapiAccuracyTestHarness): + """glm-5.2-nvfp4 / sm_103 / tep4: identity, the mtp1 variant and its acceptance.""" + + MODEL_NAME = "zai-org/GLM-5.2" + MODEL_PATH = f"{llm_models_root()}/GLM-5.2-NVFP4" + + # The parallel topology the path's tep4 segment declares: attention heads + # and experts both sharded four ways, attention DP off. Routing derives the + # target *from* these; attention DP would be a different weight layout and + # so a different target. + TEP4 = dict(tensor_parallel_size=4, moe_expert_parallel_size=4, enable_attention_dp=False) + + # The engine bounds the target's kernels are compiled for. Routing does not + # read them, so they are set here rather than left to the defaults: + # max_batch_size 32 the CUDA-graph capture ceiling the target was + # validated at. The sparse-MLA decode kernel is + # compiled for the engine's max_batch_size. + # max_num_tokens 8192 the routed-expert kernels are warmed for, and + # chunk a step at, 8192 tokens. + # max_seq_len 8192 the decode indexer's column count is fixed from it, + # and GSM8K never comes near it. + BOUNDS = dict( + max_batch_size=32, + max_num_tokens=8192, + max_seq_len=8192, + cuda_graph_config=CudaGraphConfig(max_batch_size=32, enable_padding=True), + ) + + # The checkpoint carries one MTP layer, so a draft length of 1 resolves to + # the one-model MTP-Eagle loop, the only draft path the target implements; + # its specialised rope kernel is built for that two-token width. + MTP1 = MTPDecodingConfig(max_draft_len=1) + + # The kv-cache fractions are run requirements rather than tuning choices. + # The MTP module and its separate draft pool take several GiB per rank, and + # an MTP engine on this checkpoint at the identity leg's 0.7 has been seen + # to stall rather than finish. + IDENTITY_KV = KvCacheConfig(free_gpu_memory_fraction=0.7) + MTP1_KV = KvCacheConfig(free_gpu_memory_fraction=0.55) + + # The acceptance anchor. It names the target and variant, not a test + # function, because that is what the number is a property of. Its value was + # measured on stock modeling (modeling_v2="off") under the same + # speculative_config, so the gate is not read against the target's own number. + ACCEPTANCE_KEY = "ModelingV2Glm52Nvfp4Sm103Tep4::mtp1" + + # There is deliberately one test. It evaluates the identity config as its + # first leg and the mtp1 config as its second, and ``task.evaluate`` + # asserts accuracy against the reference on the way past; the mtp1 leg also + # reads its acceptance length. Separate cases would gate nothing new and + # would each cost another engine boot of a 78-layer, 4-rank model. + + @skip_not_sm103 + @pytest.mark.skip_less_device(4) + def test_gsm8k_identity_vs_mtp1_and_acceptance(self): + """The identity accuracy gate, the gate on MTP not moving it, and acceptance. + + Turning MTP on must not move the answers. + + Rejection sampling holds the emitted distribution to the target + model's, so the two scores should differ only by sampling noise. + + Both legs assert accuracy against the registered reference as they + run -- this test is therefore the identity gate as well as the + comparison. + + Both measurements are taken **in this one test** on purpose: a delta + against a score recorded in some earlier session carries that + session's variance into the judgement. Paired, the variance is common + to both and cancels. + + The score is the harness default, with no ``scores_filter``. That is + the protocol the registered GLM-5.2 references were set under and the + one this target was validated with, so the anchor and the measurement + read the same metric. + + Acceptance length is the only gate that can see a miscomputed draft + layer. Rejection sampling makes a wrong draft path *slower*, not wrong: + every draft is rejected, the text stays correct, and the accuracy gates + pass. The mtp1 leg therefore records iteration stats and checks its + acceptance length against ``ACCEPTANCE_KEY``. + """ + task = GSM8K(self.MODEL_NAME) + + # One dict for both legs: the comparison only means anything if the two + # engines were built under the same mode. It also asks each engine to + # check its step contract on its first forward -- the contract is about + # metadata that engine prepared, so the only place it can be checked + # honestly is a run that prepared it; a synthesized metadata object + # would be checking the synthesizer. + modeling_v2 = dict(modeling_v2="require", env_overrides={MODELING_V2_VALIDATE_ENV: "1"}) + + with ( + scoped_env_overrides(modeling_v2["env_overrides"]), + LLM( + self.MODEL_PATH, + kv_cache_config=self.IDENTITY_KV, + **modeling_v2, + **self.TEP4, + **self.BOUNDS, + ) as llm, + ): + identity = task.evaluate(llm) + + with ( + scoped_env_overrides(modeling_v2["env_overrides"]), + LLM( + self.MODEL_PATH, + speculative_config=self.MTP1, + kv_cache_config=self.MTP1_KV, + enable_iter_perf_stats=True, + **modeling_v2, + **self.TEP4, + **self.BOUNDS, + ) as llm, + ): + mtp1 = task.evaluate(llm) + acceptance_length = compute_acceptance_length(llm) + + delta = mtp1 - identity + print(f"[modeling_v2] gsm8k identity={identity:.4f} mtp1={mtp1:.4f} delta={delta:+.4f}") + # 2 sigma at the ~0.6 stderr this benchmark reports at n=1319. + assert abs(delta) < 1.2, ( + f"MTP moved gsm8k by {delta:+.4f} (identity={identity:.4f}, " + f"mtp1={mtp1:.4f}); rejection sampling should have held the " + f"distribution, so this is not sampling noise" + ) + print(f"[AL] mtp1 acceptance_length = {acceptance_length:.3f}") + assert_acceptance_length(self.ACCEPTANCE_KEY, acceptance_length) + + _TP_SIZE = 16 _NUM_SAMPLES = 1319 _MAX_DRAFT_LEN = 7 @@ -80,9 +209,73 @@ } -def _engine_kwargs() -> dict: +@dataclass(frozen=True) +class AccuracyCase: + name: str + tp_size: int + batch_size: int + draft_len: int + max_seq_len: int + kv_max_tokens: int + al_reference: float | None + production_capacity: bool = False + sampled_probe: bool = False + + +POINT01 = AccuracyCase("tp16_b1_k7", 16, 1, 7, 16384, 32768, _AL_REFERENCE) +# Full-corpus modeling_v2=off reference for TP4/BS3/K5 on SM107 (2026-10-28), +# using the same 95% acceptance-length retention policy as POINT01. +POINT02 = AccuracyCase( + "tp4_b3_k5", + 4, + 3, + 5, + 1 << 20, + (1 << 20) + 64 * 3, + 4.53227789271686, + production_capacity=True, + sampled_probe=True, +) + + +# Full-corpus modeling_v2=off reference for TP8/BS2/K3 on SM107 (2026-10-28). +POINT03 = AccuracyCase( + "tp8_b2_k3", + 8, + 2, + 3, + 1 << 20, + (1 << 20) + 64 * 2, + 3.539837056750006, + production_capacity=True, + sampled_probe=True, +) + + +POINT04 = AccuracyCase( + "tp8_b1_k7", + 8, + 1, + 7, + 1 << 20, + (1 << 20) + 64, + 5.100986127360856, + production_capacity=True, + sampled_probe=True, +) + + +def _rank_env(case: AccuracyCase) -> dict[str, str]: + env = dict(_RANK_ENV) + if case in (POINT02, POINT03, POINT04): + env["TRTLLM_MOE_FINALIZE_CTA1"] = "1" + return env + + +def _engine_kwargs(case: AccuracyCase = POINT01, *, modeling_v2: str = "require") -> dict: + assert modeling_v2 in ("require", "off") speculative_config = MTPDecodingConfig( - max_draft_len=_MAX_DRAFT_LEN, + max_draft_len=case.draft_len, use_relaxed_acceptance_for_thinking=False, ) # This is a Pydantic private attribute: pass the object, not model_dump(). @@ -90,13 +283,13 @@ def _engine_kwargs() -> dict: # draft cache. Its CUDA kernels in the probe also verify this on the workers. speculative_config._allow_separate_draft_kv_cache = False return dict( - tensor_parallel_size=_TP_SIZE, + tensor_parallel_size=case.tp_size, pipeline_parallel_size=1, moe_expert_parallel_size=1, enable_attention_dp=False, - modeling_v2="require", - max_batch_size=1, - max_seq_len=16384, + modeling_v2=modeling_v2, + max_batch_size=case.batch_size, + max_seq_len=case.max_seq_len, max_num_tokens=4096, enable_chunked_prefill=True, kv_cache_config=KvCacheConfig( @@ -104,10 +297,12 @@ def _engine_kwargs() -> dict: tokens_per_block=64, enable_block_reuse=False, free_gpu_memory_fraction=0.8, - max_tokens=32768, + max_tokens=case.kv_max_tokens, host_cache_size=0, ), - cuda_graph_config=CudaGraphConfig(batch_sizes=[1], enable_padding=True), + cuda_graph_config=CudaGraphConfig( + batch_sizes=list(range(1, case.batch_size + 1)), enable_padding=True + ), moe_config=MoeConfig(backend="TRTLLM", use_low_precision_moe_combine=True), sparse_attention_config=DeepSeekSparseAttentionConfig( enable_heuristic_topk=True, @@ -117,19 +312,20 @@ def _engine_kwargs() -> dict: enable_iter_perf_stats=True, iter_perf_stats_interval=1, max_stats_len=-1, - env_overrides=dict(_RANK_ENV), + env_overrides=_rank_env(case), ) -def _assert_optimized_paths(llm: LLM, trace_dir: Path) -> None: +def _assert_optimized_paths(llm: LLM, trace_dir: Path, *, optimized: bool = True) -> None: """Observe GPU execution, which Python forward counters miss on graph replay.""" trace_dir.mkdir(parents=True) - llm.start_profile( - output_dir=str(trace_dir), - num_steps=4, - start_step=1, - activities=["GPU"], - ) + if optimized: + llm.start_profile( + output_dir=str(trace_dir), + num_steps=4, + start_step=1, + activities=["GPU"], + ) try: output = llm.generate( "Explain step by step how to multiply two three-digit numbers.", @@ -142,7 +338,11 @@ def _assert_optimized_paths(llm: LLM, trace_dir: Path) -> None: ) assert len(output.outputs[0].token_ids) == 64 finally: - llm.stop_profile() + if optimized: + llm.stop_profile() + + if not optimized: + return # These kernels are exclusive to the GLM optimized decode and MTP paths. # Requiring every TP rank prevents a route-selection or warmup-only pass. @@ -189,18 +389,293 @@ def _full_accuracy_scores(result_path: Path) -> dict[str, float]: return {name: scores[f"exact_match,{name}"] for name in ("strict-match", "flexible-extract")} -class TestModelingV2Glm52Nvfp4Sm107Tp16(LlmapiAccuracyTestHarness): - @skip_no_rubin - @pytest.mark.skip_less_mpi_world_size(_TP_SIZE) - def test_mtp7_gsm8k_and_acceptance(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path): - # The generic accuracy harness's one-sample integration mode cannot - # satisfy this full-corpus guard. - assert os.environ.get("INTEGRATION_TEST") != "1", "This guard requires all 1319 samples" - monkeypatch.setenv("TLLM_EVAL_MAX_IN_FLIGHT", "1") +def _write_json(path: Path, value) -> None: + path.write_text(json.dumps(value, indent=2) + "\n") + + +def _speculative_stats(stats: list[dict]) -> list[dict]: + return [ + stat["specDecodingStats"] + for stat in stats + if stat.get("specDecodingStats") and stat["specDecodingStats"]["numDraftTokens"] > 0 + ] + + +def _assert_point02_trace(trace_dir: Path, batch_size: int, *, sampled: bool) -> None: + traces = list(trace_dir.glob("trtllm-trace-*-rank-*.json")) + assert len(traces) == POINT02.tp_size, traces + router = { + 1: "routingFusedRouterKernel<6>", + 2: "routingFusedRouterWideKernel<12>", + 3: "routingFusedRouterMmaKernel<18>", + }[batch_size] + required = [ + router, + "glmSparseMlaFlashDecodeHeads16Kernel", + "FP8MQALogitsKernel", + "GvrMainKernel", + "draftSampleGlueWideKernel", + "_stash_last_accepted_rows_kernel", + ] + if batch_size == 1: + required.append("lowMBf16MoeFc2Kernel") + else: + required.extend(["lowMBf16MoeMultiFc1Kernel", "lowMBf16MoeMultiFc2Kernel"]) + forbidden = ("dsaHistTopKKernel", "dsaHistShardMetaKernel", "draftEhprojTableKernel") + evidence = {} + for rank in range(POINT02.tp_size): + files = [path for path in traces if path.name.endswith(f"-rank-{rank}.json")] + assert len(files) == 1, f"Missing or duplicate trace for rank {rank}" + names = sorted( + { + event["name"] + for event in json.loads(files[0].read_text())["traceEvents"] + if event.get("cat") == "kernel" + } + ) + # Profiler demangling may add whitespace inside template arguments. + compact = [re.sub(r"\s+", "", name) for name in names] + for kernel in required: + assert any(kernel in name for name in compact), ( + f"Rank {rank}, batch {batch_size}: missing {kernel}" + ) + for kernel in forbidden: + assert not any(kernel in name for name in names), ( + f"Rank {rank}, batch {batch_size}: unexpected {kernel}" + ) + assert any("glm_moe_dsa" in name and "reduceKernel" in name for name in names), ( + f"Rank {rank}: owned IPC reduction did not execute" + ) + if sampled: + for kernel in ("directPartialsKernel", "directDrawKernel"): + assert any("glm_moe_dsa" in name and kernel in name for name in names), ( + f"Rank {rank}: owned non-greedy {kernel} did not execute" + ) + evidence[str(rank)] = names + _write_json(trace_dir / "kernel-evidence.json", evidence) + + +def _assert_tp8_trace( + trace_dir: Path, case: AccuracyCase, batch_size: int, *, sampled: bool +) -> None: + assert case in (POINT03, POINT04), case + traces = list(trace_dir.glob("trtllm-trace-*-rank-*.json")) + assert len(traces) == case.tp_size, traces + required = [ + f"routingFusedRouterKernel<{batch_size * (case.draft_len + 1)}>", + "glmSparseMlaFlashDecodeKernel<8,", + "FP8MQALogitsKernel", + "GvrMainKernel", + "gemm2K256Kernel", + "mnnvlHopKernel<8,true>", + "mnnvlMoeFinalizeHopKernel<8>", + "draftEhprojTableKernel<8>", + "draftSampleGlueKernel", + "_stash_last_accepted_rows_kernel", + ] + if batch_size == 1: + required.append("lowMBf16MoeFc2Kernel") + else: + required.extend(["lowMBf16MoeMultiFc1Kernel", "lowMBf16MoeMultiFc2Kernel"]) + if sampled: + required.extend(["directPartialsKernel", "directDrawKernel"]) + forbidden = ( + "dsaHistTopKKernel", + "dsaHistShardMetaKernel", + "glmSparseMlaFlashDecodeHeads16Kernel", + "draftSampleGlueWideKernel", + ) + evidence = {} + for rank in range(case.tp_size): + files = [path for path in traces if path.name.endswith(f"-rank-{rank}.json")] + assert len(files) == 1, f"Missing or duplicate trace for rank {rank}" + names = sorted( + { + event["name"] + for event in json.loads(files[0].read_text())["traceEvents"] + if event.get("cat") == "kernel" + } + ) + compact = [re.sub(r"\s+", "", name) for name in names] + for kernel in required: + assert any(kernel in name for name in compact), ( + f"Rank {rank}, batch {batch_size}: missing {kernel}" + ) + for kernel in forbidden: + assert not any(kernel in name for name in compact), ( + f"Rank {rank}, batch {batch_size}: unexpected {kernel}" + ) + evidence[str(rank)] = names + _write_json(trace_dir / "kernel-evidence.json", evidence) + + +def _batch_probe( + llm: LLM, + case: AccuracyCase, + output_dir: Path, + batch_size: int, + *, + optimized: bool, + sampled: bool = False, +) -> list[dict]: + label = f"probe-b{batch_size}-{'sampled' if sampled else 'greedy'}" + probe_dir = output_dir / label + probe_dir.mkdir() + # >2048 keys forces real indexer scoring/TopK rather than the short-context + # dense bypass. Equal-length pretokenized prompts keep the requested batch resident. + prompts = [] + for row in range(batch_size): + tokens = llm.tokenizer.encode( + f"Example {row + 1}. " + "Explain multiplication and show each arithmetic step. " * 1024 + ) + assert len(tokens) >= 3072 + prompts.append(list(tokens[:3072])) + assert len({tuple(prompt) for prompt in prompts}) == batch_size + trace_dir = probe_dir / "traces" + if optimized: + trace_dir.mkdir() + # Skip chunked-prefill iterations, then observe several real decode + # iterations including captured graphs, not initialization warmup. + llm.start_profile(output_dir=str(trace_dir), num_steps=8, start_step=8, activities=["GPU"]) + try: + outputs = llm.generate( + prompts, + SamplingParams( + temperature=1.0 if sampled else 0.0, + top_p=0.95 if sampled else 1.0, + max_tokens=128, + ignore_eos=True, + ), + ) + finally: + if optimized: + llm.stop_profile() + stats = _request_stats(llm, batch_size) + _write_json(probe_dir / "iteration-stats.json", stats) + _write_json(probe_dir / "outputs.json", [list(out.outputs[0].token_ids) for out in outputs]) + assert len(outputs) == batch_size + assert all(len(out.outputs[0].token_ids) == 128 for out in outputs) + counts = Counter(stat["numRequestsWithDraftTokens"] for stat in _speculative_stats(stats)) + _write_json(probe_dir / "batch-counts.json", dict(counts)) + assert counts[batch_size] > 0, f"Requested batch {batch_size} was never executed: {counts}" + if optimized: + if case in (POINT03, POINT04): + _assert_tp8_trace(trace_dir, case, batch_size, sampled=sampled) + else: + _assert_point02_trace(trace_dir, batch_size, sampled=sampled) + return stats + + +def _capacity_evidence(llm: LLM, stats: list[dict], output_dir: Path, case: AccuracyCase) -> None: + """Measure final leader state; other pure-TP ranks are inferred, not inspected.""" + capacity = llm.get_kv_cache_capacity() + snapshots = [ + { + key: stat[key] + for key in ( + "kvCacheStats", + "kvCacheIterationStats", + "kvCacheIterationStatsByPoolGroup", + "kvCacheIterationStatsByLifecycle", + ) + if key in stat + } + for stat in stats + ] + snapshots = [snapshot for snapshot in snapshots if snapshot] + windows = { + int(lifecycle["windowSize"]) + for snapshot in snapshots + for lifecycle in snapshot.get("kvCacheIterationStatsByLifecycle", {}).values() + if lifecycle.get("kind") == "attention" and lifecycle.get("windowSize") is not None + } + evidence = dict( + measured_rank=0, + capacity=capacity, + attention_windows=sorted(windows), + snapshots=snapshots, + other_rank_capacity="inferred from V2 distributed MIN and replicated pure-TP cache schema", + route_evidence="Every TP rank is separately checked by GPU traces; traces do not measure capacity", + ) + _write_json(output_dir / "capacity-evidence.json", evidence) + assert capacity.get("tokensPerBlock") == 64 and capacity.get("maxNumBlocks", 0) > 0, capacity + # GLM uses full-attention lifecycles. Their V2 windowSize is the final + # manager.max_seq_len after profiling/capacity clamping, not llm.args. + assert capacity.get("maxNumTokens", 0) >= case.max_seq_len, capacity + assert windows and min(windows) >= case.max_seq_len, ( + f"Actual leader window below {case.max_seq_len}: {windows}" + ) + assert any(snapshot.get("kvCacheIterationStatsByPoolGroup") for snapshot in snapshots), ( + "Missing physical pool-group capacity statistics" + ) + + +def run_accuracy_case( + case: AccuracyCase, output_dir: Path, *, modeling_v2: str = "require" +) -> dict: + """Shared protocol; a task-local stock runner may explicitly select ``off``. + + Repository tests always require ModelingV2 and a frozen independent AL + reference. Stock records the reference without a candidate retention check. + """ + assert modeling_v2 in ("require", "off") + assert os.environ.get("INTEGRATION_TEST") != "1", "This guard requires all 1319 samples" + if modeling_v2 == "require": + assert case.al_reference is not None, ( + f"{case.name}: run and review the matched full1319 modeling_v2=off baseline, " + "then freeze its independent AL reference before running the repository target test" + ) + output_dir.mkdir(parents=True, exist_ok=True) + _write_json( + output_dir / "accuracy-protocol.json", + dict( + case=asdict(case), + modeling_v2=modeling_v2, + num_samples=_NUM_SAMPLES, + model_path=_MODEL_PATH, + dataset_path=_DATASET_PATH, + test_source_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), + scoring_temperature=0.0, + scoring_top_p=1.0, + max_output_tokens=12288, + max_prompt_tokens=4096, + forced_accepted_tokens=0, + rank_env=_rank_env(case), + ), + ) + with pytest.MonkeyPatch.context() as monkeypatch: + monkeypatch.setenv("TLLM_EVAL_MAX_IN_FLIGHT", str(case.batch_size)) monkeypatch.setenv("TLLM_EVAL_PARTIAL_SCORES_EVERY", "100") - with scoped_env_overrides(_RANK_ENV), LLM(_MODEL_PATH, **_engine_kwargs()) as llm: - _assert_optimized_paths(llm, tmp_path / "optimized-path-traces") - _request_stats(llm, 1) # Exclude the coverage probe from corpus AL. + with ( + scoped_env_overrides(_rank_env(case)), + LLM(_MODEL_PATH, **_engine_kwargs(case, modeling_v2=modeling_v2)) as llm, + ): + if case.production_capacity: + probe_stats = [] + probe_batches = {POINT02: (3, 1, 2), POINT03: (2, 1), POINT04: (1,)} + for batch in probe_batches[case]: + probe_stats.extend( + _batch_probe( + llm, case, output_dir, batch, optimized=modeling_v2 == "require" + ) + ) + if case.sampled_probe: + probe_stats.extend( + _batch_probe( + llm, + case, + output_dir, + case.batch_size, + optimized=modeling_v2 == "require", + sampled=True, + ) + ) + _capacity_evidence(llm, probe_stats, output_dir, case) + else: + _assert_optimized_paths( + llm, output_dir / "optimized-path-traces", optimized=modeling_v2 == "require" + ) + _write_json(output_dir / "probe-iteration-stats.json", _request_stats(llm, 1)) evaluator = GSM8KInferenceX( dataset_path=_DATASET_PATH, num_samples=_NUM_SAMPLES, @@ -211,186 +686,72 @@ def test_mtp7_gsm8k_and_acceptance(self, monkeypatch: pytest.MonkeyPatch, tmp_pa shuffle_dataset=False, fewshot_random_seed=1234, log_samples=True, - output_path=str(tmp_path), + output_path=str(output_dir), ) evaluator.evaluate( llm, sampling_params=SamplingParams( - temperature=0.0, - top_p=1.0, - max_tokens=12288, - truncate_prompt_tokens=4096, + temperature=0.0, top_p=1.0, max_tokens=12288, truncate_prompt_tokens=4096 ), streaming=False, sampling_override=True, scores_filter="exact_match,strict-match", ) stats = _request_stats(llm, _NUM_SAMPLES) - - spec_stats = [ - stat["specDecodingStats"] - for stat in stats - if stat.get("specDecodingStats") and stat["specDecodingStats"]["numDraftTokens"] > 0 - ] - assert len(spec_stats) >= _NUM_SAMPLES, "Missing full-run speculative iteration statistics" - assert all(stat["numRequestsWithDraftTokens"] == 1 for stat in spec_stats) - acceptance_length = acceptance_length_from_iteration_stats(stats) - scores = _full_accuracy_scores(tmp_path / f"samples_{_TASK_NAME}.json") - metrics = dict(scores) - metrics["acceptance_length"] = acceptance_length - metrics["speculative_iterations"] = len(spec_stats) - metrics["accepted_draft_tokens"] = sum(stat["numAcceptedTokens"] for stat in spec_stats) - metrics["drafted_tokens"] = sum(stat["numDraftTokens"] for stat in spec_stats) - metrics["speculative_request_steps"] = sum( - stat["numRequestsWithDraftTokens"] for stat in spec_stats - ) - metrics["completed_requests"] = sum(stat["numCompletedRequests"] for stat in stats) - metrics["new_requests"] = sum(stat["numNewActiveRequests"] for stat in stats) - (tmp_path / "glm_metrics.json").write_text(json.dumps(metrics, indent=2) + "\n") - print(f"[GLM TP16 c1] {json.dumps(metrics)}") - for name, score in scores.items(): - assert _ACCURACY_FLOOR <= score <= 100.0, ( - f"{name}: {score:.3f}% is below the local {_ACCURACY_FLOOR}% floor" - ) - assert _AL_FLOOR <= acceptance_length <= _MAX_DRAFT_LEN + 1, ( - f"MTP acceptance length {acceptance_length:.3f} is outside " - f"[{_AL_FLOOR:.3f}, {_MAX_DRAFT_LEN + 1}]" + _write_json(output_dir / "corpus-iteration-stats.json", stats) + spec_stats = _speculative_stats(stats) + counts = Counter(stat["numRequestsWithDraftTokens"] for stat in spec_stats) + assert len(spec_stats) >= (_NUM_SAMPLES + case.batch_size - 1) // case.batch_size, ( + "Missing full-run speculative iteration statistics" + ) + assert all(1 <= count <= case.batch_size for count in counts), counts + assert counts[case.batch_size] > 0, f"Corpus never executed the full batch: {counts}" + acceptance_length = acceptance_length_from_iteration_stats(stats) + scores = _full_accuracy_scores(output_dir / f"samples_{_TASK_NAME}.json") + metrics = dict(scores) + metrics.update( + acceptance_length=acceptance_length, + speculative_iterations=len(spec_stats), + accepted_draft_tokens=sum(stat["numAcceptedTokens"] for stat in spec_stats), + drafted_tokens=sum(stat["numDraftTokens"] for stat in spec_stats), + speculative_request_steps=sum(stat["numRequestsWithDraftTokens"] for stat in spec_stats), + speculative_batch_counts=dict(counts), + completed_requests=sum(stat["numCompletedRequests"] for stat in stats), + new_requests=sum(stat["numNewActiveRequests"] for stat in stats), + modeling_v2=modeling_v2, + al_reference=case.al_reference, + ) + _write_json(output_dir / "glm_metrics.json", metrics) + print(f"[GLM {case.name} {modeling_v2}] {json.dumps(metrics)}") + assert metrics["completed_requests"] == metrics["new_requests"] == _NUM_SAMPLES, metrics + for name, score in scores.items(): + assert _ACCURACY_FLOOR <= score <= 100.0, ( + f"{name}: {score:.3f}% is below the local {_ACCURACY_FLOOR}% floor" ) - - -# The targets assert their own SM at construction: certification is per GPU -# architecture, and a receipt from another one says nothing here. -skip_not_sm103 = pytest.mark.skipif( - get_sm_version() != 103, reason="modeling_v2 targets in this batch are certified on sm_103 only" -) - - -class TestModelingV2Glm52Nvfp4Sm103Tep4(LlmapiAccuracyTestHarness): - """glm-5.2-nvfp4 / sm_103 / tep4: identity, the mtp1 variant and its acceptance.""" - - MODEL_NAME = "zai-org/GLM-5.2" - MODEL_PATH = f"{llm_models_root()}/GLM-5.2-NVFP4" - - # The parallel topology the path's tep4 segment declares: attention heads - # and experts both sharded four ways, attention DP off. Routing derives the - # target *from* these; attention DP would be a different weight layout and - # so a different target. - TEP4 = dict(tensor_parallel_size=4, moe_expert_parallel_size=4, enable_attention_dp=False) - - # The engine bounds the target's kernels are compiled for. Routing does not - # read them, so they are set here rather than left to the defaults: - # max_batch_size 32 the CUDA-graph capture ceiling the target was - # validated at. The sparse-MLA decode kernel is - # compiled for the engine's max_batch_size. - # max_num_tokens 8192 the routed-expert kernels are warmed for, and - # chunk a step at, 8192 tokens. - # max_seq_len 8192 the decode indexer's column count is fixed from it, - # and GSM8K never comes near it. - BOUNDS = dict( - max_batch_size=32, - max_num_tokens=8192, - max_seq_len=8192, - cuda_graph_config=CudaGraphConfig(max_batch_size=32, enable_padding=True), + floor = 0.95 * case.al_reference if modeling_v2 == "require" else 1.0 + assert floor <= acceptance_length <= case.draft_len + 1, ( + f"MTP AL {acceptance_length:.6f} is outside [{floor:.6f}, {case.draft_len + 1}]" ) + return metrics - # The checkpoint carries one MTP layer, so a draft length of 1 resolves to - # the one-model MTP-Eagle loop, the only draft path the target implements; - # its specialised rope kernel is built for that two-token width. - MTP1 = MTPDecodingConfig(max_draft_len=1) - # The kv-cache fractions are run requirements rather than tuning choices. - # The MTP module and its separate draft pool take several GiB per rank, and - # an MTP engine on this checkpoint at the identity leg's 0.7 has been seen - # to stall rather than finish. - IDENTITY_KV = KvCacheConfig(free_gpu_memory_fraction=0.7) - MTP1_KV = KvCacheConfig(free_gpu_memory_fraction=0.55) - - # The acceptance anchor. It names the target and variant, not a test - # function, because that is what the number is a property of. Its value was - # measured on stock modeling (modeling_v2="off") under the same - # speculative_config, so the gate is not read against the target's own number. - ACCEPTANCE_KEY = "ModelingV2Glm52Nvfp4Sm103Tep4::mtp1" - - # There is deliberately one test. It evaluates the identity config as its - # first leg and the mtp1 config as its second, and ``task.evaluate`` - # asserts accuracy against the reference on the way past; the mtp1 leg also - # reads its acceptance length. Separate cases would gate nothing new and - # would each cost another engine boot of a 78-layer, 4-rank model. - - @skip_not_sm103 - @pytest.mark.skip_less_device(4) - def test_gsm8k_identity_vs_mtp1_and_acceptance(self): - """The identity accuracy gate, the gate on MTP not moving it, and acceptance. - - Turning MTP on must not move the answers. - - Rejection sampling holds the emitted distribution to the target - model's, so the two scores should differ only by sampling noise. - - Both legs assert accuracy against the registered reference as they - run -- this test is therefore the identity gate as well as the - comparison. - - Both measurements are taken **in this one test** on purpose: a delta - against a score recorded in some earlier session carries that - session's variance into the judgement. Paired, the variance is common - to both and cancels. - - The score is the harness default, with no ``scores_filter``. That is - the protocol the registered GLM-5.2 references were set under and the - one this target was validated with, so the anchor and the measurement - read the same metric. - - Acceptance length is the only gate that can see a miscomputed draft - layer. Rejection sampling makes a wrong draft path *slower*, not wrong: - every draft is rejected, the text stays correct, and the accuracy gates - pass. The mtp1 leg therefore records iteration stats and checks its - acceptance length against ``ACCEPTANCE_KEY``. - """ - task = GSM8K(self.MODEL_NAME) - - # One dict for both legs: the comparison only means anything if the two - # engines were built under the same mode. It also asks each engine to - # check its step contract on its first forward -- the contract is about - # metadata that engine prepared, so the only place it can be checked - # honestly is a run that prepared it; a synthesized metadata object - # would be checking the synthesizer. - modeling_v2 = dict(modeling_v2="require", env_overrides={MODELING_V2_VALIDATE_ENV: "1"}) +class TestModelingV2Glm52Nvfp4Sm107LowLatency(LlmapiAccuracyTestHarness): + @skip_no_rubin + @pytest.mark.skip_less_mpi_world_size(16) + def test_mtp7_gsm8k_and_acceptance(self, tmp_path: Path): + run_accuracy_case(POINT01, tmp_path) - with ( - scoped_env_overrides(modeling_v2["env_overrides"]), - LLM( - self.MODEL_PATH, - kv_cache_config=self.IDENTITY_KV, - **modeling_v2, - **self.TEP4, - **self.BOUNDS, - ) as llm, - ): - identity = task.evaluate(llm) + @skip_no_rubin + @pytest.mark.skip_less_mpi_world_size(4) + def test_tp4_mtp5_gsm8k_and_acceptance(self, tmp_path: Path): + run_accuracy_case(POINT02, tmp_path) - with ( - scoped_env_overrides(modeling_v2["env_overrides"]), - LLM( - self.MODEL_PATH, - speculative_config=self.MTP1, - kv_cache_config=self.MTP1_KV, - enable_iter_perf_stats=True, - **modeling_v2, - **self.TEP4, - **self.BOUNDS, - ) as llm, - ): - mtp1 = task.evaluate(llm) - acceptance_length = compute_acceptance_length(llm) + @skip_no_rubin + @pytest.mark.skip_less_mpi_world_size(8) + def test_tp8_mtp3_gsm8k_and_acceptance(self, tmp_path: Path): + run_accuracy_case(POINT03, tmp_path) - delta = mtp1 - identity - print(f"[modeling_v2] gsm8k identity={identity:.4f} mtp1={mtp1:.4f} delta={delta:+.4f}") - # 2 sigma at the ~0.6 stderr this benchmark reports at n=1319. - assert abs(delta) < 1.2, ( - f"MTP moved gsm8k by {delta:+.4f} (identity={identity:.4f}, " - f"mtp1={mtp1:.4f}); rejection sampling should have held the " - f"distribution, so this is not sampling noise" - ) - print(f"[AL] mtp1 acceptance_length = {acceptance_length:.3f}") - assert_acceptance_length(self.ACCEPTANCE_KEY, acceptance_length) + @skip_no_rubin + @pytest.mark.skip_less_mpi_world_size(8) + def test_tp8_mtp7_gsm8k_and_acceptance(self, tmp_path: Path) -> None: + run_accuracy_case(POINT04, tmp_path) diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 31ba50c6488d..fcc79b6af1db 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -55,7 +55,7 @@ l0_cpu: - unittest/_torch/compilation/test_remove_copy_pass.py::test_remove_copy_preserves_minimax_producer_outputs - unittest/_torch/distributed # GLM routing and draft workspace setup use CPU metadata and mocked allocations. - - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_exchange.py + - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_draft_exchange.py - unittest/_torch/modeling_v2/models/glm_moe_dsa/test_routing.py - unittest/_torch/pyexecutor - unittest/_torch/test_connector.py::test_a_partial_grouped_scheduler_is_rejected_at_bring_up diff --git a/tests/integration/test_lists/test-db/l0_gr100.yml b/tests/integration/test_lists/test-db/l0_gr100.yml index a9cbbc03f821..310fed2d3bab 100644 --- a/tests/integration/test_lists/test-db/l0_gr100.yml +++ b/tests/integration/test_lists/test-db/l0_gr100.yml @@ -19,10 +19,10 @@ l0_gr100: - unittest/_torch/modules/kimi_kda/test_kda_flashinfer_prefill.py::test_auto_prefill_stays_finite_on_repeated_token_with_mixed_decay # The GLM-5.2 target's unit tests, one entry per op directory. Its comm tests # run 4 ranks, and CI has no 4-GPU GR100 stage yet. - - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention - - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa - - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/embedding - - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm - - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe - - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/quantization - - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec + - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/attention + - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa + - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/embedding + - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm + - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/moe + - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/quantization + - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index f8419381b754..9d0db4b51516 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -152,9 +152,6 @@ unittest/_torch/kernels/parallel/test_fp4_linear.py::test_fp4_gemm_bias_per_back unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/comm SKIP (https://nvbugs/7011673) unittest/_torch/modeling_v2/comm/test_modeling_v2_allgather_op_matrix.py::test_allgather_op_matrix SKIP (https://nvbugs/6843474) unittest/_torch/modeling_v2/comm/test_modeling_v2_reducescatter_op_matrix.py::test_reducescatter_op_matrix SKIP (https://nvbugs/6843474) -unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_topk_op_matrix.py::test_hist_topk_op_matrix SKIP (https://nvbugs/7004363) -unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_ehproj_table.py::test_draft_ehproj_table_op_matrix SKIP (https://nvbugs/7004363) -unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_sample_glue.py::test_draft_sample_glue_op_matrix SKIP (https://nvbugs/7004363) unittest/_torch/modeling_v2/test_modeling_v2_env_scope.py SKIP (https://nvbugs/7024242) unittest/_torch/modules/mamba/test_flashinfer_gdn_decode.py::test_fi_decode_misaligned_index_slice SKIP (https://nvbugs/6842831) unittest/_torch/modules/tests_lora_modules/test_nemotron_h_lora_sanity.py::TestNemotronHLoRA::test_lora_pp2_sanity SKIP (https://nvbugs/6428124) diff --git a/tests/unittest/_torch/kernels/parallel_hw_agnostic/test_custom_ops.py b/tests/unittest/_torch/kernels/parallel_hw_agnostic/test_custom_ops.py index c39ee3ada042..81dc9a42e6f3 100644 --- a/tests/unittest/_torch/kernels/parallel_hw_agnostic/test_custom_ops.py +++ b/tests/unittest/_torch/kernels/parallel_hw_agnostic/test_custom_ops.py @@ -22,7 +22,7 @@ # The MiniMax-M3 ops' fake impls are registered by their modeling_v2 package. import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.fake_ops # noqa: F401 import tensorrt_llm._torch.pyexecutor.step_inputs # noqa: F401 -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops import ( +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops import ( # noqa: E501 _cpp_fakes, # noqa: F401 ) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_mla_kv_half_append.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/attention/test_modeling_v2_mla_kv_half_append.py similarity index 98% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_mla_kv_half_append.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/attention/test_modeling_v2_mla_kv_half_append.py index c3ad37ea9d02..7fbb0b2da5fb 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_mla_kv_half_append.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/attention/test_modeling_v2_mla_kv_half_append.py @@ -15,12 +15,12 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.attention.mla_kv_half_append import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.attention.mla_kv_half_append import ( # noqa: E501 KV_LORA_RANK, LATENT_OFFSET, ROPE_DIM, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.attention.mla_kv_half_append import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.attention.mla_kv_half_append import ( # noqa: E501 mla_kv_half_append as op, ) from tensorrt_llm._torch.attention.backends.interface import RopeParams diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_sparse_mla_flash_decode.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/attention/test_modeling_v2_sparse_mla_flash_decode.py similarity index 72% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_sparse_mla_flash_decode.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/attention/test_modeling_v2_sparse_mla_flash_decode.py index a255b9e9819b..a8c4b61b1fba 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention/test_modeling_v2_sparse_mla_flash_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/attention/test_modeling_v2_sparse_mla_flash_decode.py @@ -21,24 +21,34 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.attention.mla_kv_half_append import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import true_fp32_matmul +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.attention.mla_kv_half_append import ( # noqa: E501 LATENT_OFFSET, mla_kv_half_append, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.attention.sparse_mla_flash_decode import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.attention.sparse_mla_flash_decode import ( # noqa: E501 GLM_SOFTMAX_SCALE, HEAD_DIM_QK, KV_LORA_RANK, + LATENT_NUM_HEADS, MAX_TOP_K, NUM_HEADS, V_HEAD_DIM, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.attention.sparse_mla_flash_decode import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.attention.sparse_mla_flash_decode import ( # noqa: E501 sparse_mla_flash_decode as op, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.attention.sparse_mla_flash_decode import ( # noqa: E501 + sparse_mla_flash_decode_latent as latent_op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.gemm._low_m_splitk import ( # noqa: E501 + bmm_prewait_tactic, + run_splitk_bmm, +) from tensorrt_llm._torch.attention.backends.interface import RopeParams from tensorrt_llm._torch.attention.backends.sparse.dsa.cache_manager import DSACacheManagerV2 from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._utils import get_sm_version from tensorrt_llm.bindings import DataType @@ -159,6 +169,7 @@ def _build(spec: dict, seed: int) -> _Case: """Turn a cell spec into the op's arguments, over a pool of the spec's geometry.""" gen_cpu = torch.Generator().manual_seed(seed) gen = torch.Generator(device="cuda").manual_seed(seed) + heads = spec.get("heads", NUM_HEADS) tokens_per_block = spec["tokens_per_block"] num_layers = spec["num_layers"] layer_offset = spec["layer_offset"] @@ -187,19 +198,19 @@ def _build(spec: dict, seed: int) -> _Case: scale = torch.tensor(GLM_SOFTMAX_SCALE, dtype=torch.float32) kwargs = dict( - q=( - torch.randn(num_tokens, NUM_HEADS, HEAD_DIM_QK, generator=gen, device="cuda") * _Q_STD - ).to(torch.float8_e4m3fn), + q=(torch.randn(num_tokens, heads, HEAD_DIM_QK, generator=gen, device="cuda") * _Q_STD).to( + torch.float8_e4m3fn + ), kv_pool=pool.view(torch.float8_e4m3fn), indices=(positions if spec["remap"] else global_rows).cuda(), seq_lens_kv=torch.tensor(kv_lens, dtype=torch.int32, device="cuda"), mla_bmm1_scale=torch.stack((scale, scale * torch.tensor(_LOG2E))).cuda(), mla_bmm2_scale=torch.ones(1, dtype=torch.float32, device="cuda"), v_b_proj=( - torch.randn(NUM_HEADS, V_HEAD_DIM, KV_LORA_RANK, generator=gen, device="cuda") * _VB_STD + torch.randn(heads, V_HEAD_DIM, KV_LORA_RANK, generator=gen, device="cuda") * _VB_STD ).to(torch.bfloat16), output=torch.full( - (num_tokens + _SPARE_ROWS, NUM_HEADS * V_HEAD_DIM), + (num_tokens + _SPARE_ROWS, heads * V_HEAD_DIM), float("nan"), dtype=torch.bfloat16, device="cuda", @@ -236,7 +247,7 @@ def _unrounded_attention(case: _Case) -> torch.Tensor: weights = torch.nan_to_num(torch.softmax(scores, dim=-1), nan=0.0) attn = torch.einsum("thj,tjd->thd", weights, keys[..., :KV_LORA_RANK]) out = torch.einsum("thd,hvd->thv", attn, kw["v_b_proj"].double()) - return out.reshape(num_tokens, NUM_HEADS * V_HEAD_DIM).float() + return out.reshape(num_tokens, kw["q"].shape[1] * V_HEAD_DIM).float() def _seed(why: str) -> int: @@ -264,21 +275,28 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res assert torch.isnan(untouched).all(), "rows past the tokens were written" -def test_capture_and_replay_match_the_eager_call() -> None: - """The shared-layer call (early gather) captured in a CUDA graph replays bit for bit, and - two eager calls agree bit for bit.""" - cell = op.CELLS[1] - assert cell.spec["gather_before_wait"] - case = _build(cell.spec, _seed(cell.why)) - output = case.kwargs["output"] - op(**case.kwargs) - torch.cuda.synchronize() - eager = output.clone() - output.fill_(float("nan")) +@pytest.mark.parametrize( + "heads,seq_len_q,kv_lens", + [(4, 8, (20011,)), (8, 4, (9000, 17003)), (8, 8, (20011,))], + ids=("tp16-b1-q8", "tp8-b2-q4", "tp8-b1-q8"), +) +def test_capture_and_replay_match_the_eager_call( + heads: int, seq_len_q: int, kv_lens: tuple[int, ...] +) -> None: + """Replay changed Q, selected indices and KV into a strided output without touching canaries.""" + spec = dict(op.CELLS[1].spec, heads=heads, seq_len_q=seq_len_q, kv_lens=kv_lens) + case = _build(spec, 7391 + heads) + columns = heads * V_HEAD_DIM + storage = torch.full( + (case.num_tokens + _SPARE_ROWS, columns + 16), + float("nan"), + dtype=torch.bfloat16, + device="cuda", + ) + output = storage[:, :columns] + case.kwargs["output"] = output op(**case.kwargs) torch.cuda.synchronize() - assert torch.equal(output[: case.num_tokens], eager[: case.num_tokens]), "not deterministic" - graph = torch.cuda.CUDAGraph() stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) @@ -286,11 +304,24 @@ def test_capture_and_replay_match_the_eager_call() -> None: with torch.cuda.graph(graph, stream=stream): op(**case.kwargs) torch.cuda.current_stream().wait_stream(stream) + gen = torch.Generator(device="cuda").manual_seed(9843 + heads) + gen_cpu = torch.Generator().manual_seed(9843 + heads) for _ in range(3): - output.fill_(float("nan")) + _fill_e4m3_(case.kwargs["q"].view(torch.uint8), gen) + _fill_e4m3_(case.kwargs["kv_pool"].view(torch.uint8), gen) + case.kwargs["indices"].copy_( + _select(kv_lens, seq_len_q, spec["top_k"], spec["holes"], gen_cpu) + ) + storage.fill_(float("nan")) graph.replay() torch.cuda.synchronize() - assert torch.equal(output[: case.num_tokens], eager[: case.num_tokens]) + replay = output[: case.num_tokens].clone() + assert torch.isnan(storage[case.num_tokens :]).all(), "spare rows were written" + assert torch.isnan(storage[:, columns:]).all(), "spare columns were written" + output.fill_(float("nan")) + op(**case.kwargs) + torch.cuda.synchronize() + assert torch.equal(output[: case.num_tokens], replay) def _guard_case(cell_index: int) -> _Case: @@ -447,7 +478,8 @@ def _convert( ) -def test_reads_the_rows_the_kv_append_writes_in_the_dsa_pool() -> None: +@pytest.mark.parametrize("heads", (4, 8), ids=("tp16-b1-q8", "tp8-b1-q8")) +def test_reads_the_rows_the_kv_append_writes_in_the_dsa_pool(heads: int) -> None: """Both ops against main's DSA cache layout. On a DSACacheManagerV2 pool, glm_mla_kv_half_append appends this step's 8 rows of one layer @@ -512,15 +544,13 @@ def test_reads_the_rows_the_kv_append_writes_in_the_dsa_pool() -> None: indices = _select(kv_lens, seq_len_q, MAX_TOP_K, False, gen_cpu).cuda() global_rows = _convert(block_table, indices, seq_len_q, page_params, tpb) - q = ( - torch.randn(num_tokens, NUM_HEADS, HEAD_DIM_QK, generator=gen, device="cuda") * _Q_STD - ).to(torch.float8_e4m3fn) + q = (torch.randn(num_tokens, heads, HEAD_DIM_QK, generator=gen, device="cuda") * _Q_STD).to( + torch.float8_e4m3fn + ) v_b_proj = ( - torch.randn(NUM_HEADS, V_HEAD_DIM, KV_LORA_RANK, generator=gen, device="cuda") * _VB_STD + torch.randn(heads, V_HEAD_DIM, KV_LORA_RANK, generator=gen, device="cuda") * _VB_STD ).to(torch.bfloat16) - output = torch.empty( - num_tokens, NUM_HEADS * V_HEAD_DIM, dtype=torch.bfloat16, device="cuda" - ) + output = torch.empty(num_tokens, heads * V_HEAD_DIM, dtype=torch.bfloat16, device="cuda") ref = op.reference( q, pool, @@ -557,3 +587,144 @@ def test_reads_the_rows_the_kv_append_writes_in_the_dsa_pool() -> None: _compare(output, ref) finally: env.shutdown() + + +def _latent_case(spec, seed): + case = _build(spec, seed) + gen = torch.Generator(device="cuda").manual_seed(seed + 1) + case.kwargs["q"] = ( + torch.randn(case.num_tokens, LATENT_NUM_HEADS, HEAD_DIM_QK, generator=gen, device="cuda") + * _Q_STD + ).to(torch.float8_e4m3fn) + del case.kwargs["v_b_proj"] + case.kwargs["output"] = torch.full( + (case.num_tokens + _SPARE_ROWS, LATENT_NUM_HEADS, KV_LORA_RANK), + float("nan"), + dtype=torch.bfloat16, + device="cuda", + ) + return case + + +@pytest.mark.parametrize("cell", latent_op.CELLS, ids=[c.why[:44] for c in latent_op.CELLS]) +@pytest.mark.parametrize("early_trigger", [False, True]) +def test_latent_certified_cells(cell, early_trigger, certifying) -> None: # noqa: F811 + certifying(latent_op) + case = _latent_case(cell.spec, _seed(cell.why)) + case.kwargs["trigger_after_wait"] = early_trigger + output = latent_op(**case.kwargs) + torch.cuda.synchronize() + _compare(output[: case.num_tokens], latent_op.reference(**case.kwargs)) + assert torch.isnan(output[case.num_tokens :].float()).all() + + +@pytest.mark.parametrize("tokens", [6, 12, 18]) +def test_latent_cuda_graph_replay(tokens) -> None: + cell = next( + c + for c in latent_op.CELLS + if c.spec["seq_len_q"] == 6 and 6 * len(c.spec["kv_lens"]) == tokens + ) + case = _latent_case(cell.spec, 9100 + tokens) + output = latent_op(**case.kwargs) + torch.cuda.synchronize() + eager = output[:tokens].clone() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + latent_op(**case.kwargs) + for _ in range(3): + output.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(output[:tokens], eager) + assert torch.isnan(output[tokens:].float()).all() + + +def test_latent_refuses_more_than_twenty_two_queries() -> None: + spec = dict(latent_op.CELLS[-1].spec, seq_len_q=23) + case = _latent_case(spec, 9201) + with pytest.raises(RuntimeError, match="at most 22"): + latent_op(**case.kwargs) + + +def test_latent_output_abi_is_explicit() -> None: + case = _latent_case(latent_op.CELLS[0].spec, 9202) + case.kwargs["output"] = torch.empty( + case.num_tokens, LATENT_NUM_HEADS * V_HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + with pytest.raises(RuntimeError, match="16, 512"): + latent_op(**case.kwargs) + + +def test_latent_padded_output_stride() -> None: + case = _latent_case(latent_op.CELLS[0].spec, 9203) + storage = torch.full( + (case.num_tokens + _SPARE_ROWS, LATENT_NUM_HEADS * KV_LORA_RANK + 8), + float("nan"), + dtype=torch.bfloat16, + device="cuda", + ) + case.kwargs["output"] = storage[:, : LATENT_NUM_HEADS * KV_LORA_RANK].view( + case.num_tokens + _SPARE_ROWS, LATENT_NUM_HEADS, KV_LORA_RANK + ) + latent_op(**case.kwargs) + torch.cuda.synchronize() + _compare(case.kwargs["output"][: case.num_tokens], latent_op.reference(**case.kwargs)) + assert torch.isnan(storage[:, -8:].float()).all() + assert torch.isnan(storage[case.num_tokens :].float()).all() + + +@pytest.mark.parametrize("tokens", [1, 3, 6, 12, 18]) +def test_latent_early_trigger_with_prewait_v_b(tokens) -> None: + """The consumer preloads weights, waits for latent output, and writes token-major v_b rows.""" + cell = next( + c for c in latent_op.CELLS if c.spec["seq_len_q"] * len(c.spec["kv_lens"]) == tokens + ) + case = _latent_case(cell.spec, 9300 + tokens) + case.kwargs["trigger_after_wait"] = True + gen = torch.Generator(device="cuda").manual_seed(9400 + tokens) + weights = ( + torch.randn(LATENT_NUM_HEADS, V_HEAD_DIM, KV_LORA_RANK, generator=gen, device="cuda") + * _VB_STD + ).to(torch.bfloat16) + projected = torch.empty( + tokens, LATENT_NUM_HEADS, V_HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + tactic = bmm_prewait_tactic(tokens, V_HEAD_DIM, KV_LORA_RANK) + assert tactic is not None + + def step(): + latent_op(**case.kwargs) + run_splitk_bmm( + case.kwargs["output"][:tokens].transpose(0, 1), + weights, + projected.transpose(0, 1), + get_env_enable_pdl(), + tactic, + reserve_full_smem=False, + ) + + def reference(): + latent = latent_op.reference(**case.kwargs) + with true_fp32_matmul(): + return torch.einsum("thd,hvd->thv", latent.float(), weights.float()).to(torch.bfloat16) + + step() # Compile/load both kernels before capture. + torch.cuda.synchronize() + _compare(projected, reference()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + step() + q = case.kwargs["q"] + for seed in (9500 + tokens, 9600 + tokens): + gen.manual_seed(seed) + q.copy_((torch.randn(q.shape, generator=gen, device="cuda") * _Q_STD).to(q.dtype)) + case.kwargs["output"].fill_(float("nan")) + projected.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + _compare(projected, reference()) + replay = projected.clone() + step() + torch.cuda.synchronize() + assert torch.equal(projected, replay) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/_intra_node_reduce_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/_intra_node_reduce_op_matrix.py new file mode 100644 index 000000000000..b69b47b58ace --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/_intra_node_reduce_op_matrix.py @@ -0,0 +1,349 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Four-rank IPC arithmetic, shared-workspace interleave and CUDA graph matrix. + +Run with CUDA_VISIBLE_DEVICES=0,1,2,3. The launcher kills all ranks on any failure or +at its deadline. The communication is real IPC; gloo only exchanges allocation handles. +""" + +from __future__ import annotations + +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from types import SimpleNamespace + +import torch +import torch.distributed as dist + +DEADLINE_S = 900 +_WORKER_FLAG = "--rank-worker" +HIDDEN = 6144 +EPS = 1e-5 + + +def _run_one_rank(): + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank, world = int(os.environ["RANK"]), int(os.environ["WORLD_SIZE"]) + assert world == 4 + torch.cuda.set_device(rank) + assert torch.cuda.get_device_capability() == (10, 7) + dist.init_process_group( + "cuda:nccl,cpu:gloo", init_method="env://", device_id=torch.device("cuda", rank) + ) + host_group = dist.new_group(backend="gloo") + from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm import allreduce as shared_ar + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.comm.intra_node_moe_finalize import ( # noqa: E501 + intra_node_moe_finalize, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.comm.intra_node_reduce import ( # noqa: E501 + IpcCollectives, + intra_node_reduce, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.comm.pre_moe_fp4_linear import ( # noqa: E501 + pre_moe_fp4_linear, + ) + from tensorrt_llm._torch.distributed import communicator + from tensorrt_llm._torch.distributed.ops import get_allreduce_workspace + + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating + + from tensorrt_llm._utils import torch_pybind11_abi + from tensorrt_llm.bindings.internal.process_group import init_pg + from tensorrt_llm.functional import AllReduceStrategy + from tensorrt_llm.mapping import Mapping + + class HostExchange: + def tp_allgather(self, obj, *, small_payload=False): + result = [None] * world + dist.all_gather_object(result, obj, group=host_group) + return result + + local_group = dist.new_group(ranks=list(range(world)), backend="cuda:nccl,cpu:gloo") + init_pg(dist.group.WORLD, local_group, torch_pybind11_abi()) + communicator.Distributed.get = staticmethod(lambda mapping: HostExchange()) + mapping = Mapping( + world_size=4, rank=rank, gpus_per_node=4, tp_size=4, moe_tp_size=4, moe_ep_size=1 + ) + table = get_allreduce_workspace(mapping) + ipc, reason = IpcCollectives.bind(mapping, [table, table]) + assert ipc is not None, reason + state = ipc.workspace + assert ( + state.table is table and state.flags is state.owners[3] and state.layout is state.owners[4] + ) + # One rank cannot bind a lookalike table: all ranks must decline before a launch. + rejected, reason = IpcCollectives.bind(mapping, [table.clone() if rank == 0 else table]) + assert rejected is None and "rank 0" in reason + assert not ipc.supported(25) and not ipc.supported(0) + args = shared_ar.allreduce_group_args(mapping, int(AllReduceStrategy.ONESHOT)) + assert args["workspace"] is table + + def tensor(rows, seed, r=rank): + g = torch.Generator(device="cuda").manual_seed(seed * 1009 + r) + return torch.randn((rows, HIDDEN), device="cuda", generator=g).to(torch.bfloat16) + + def norm_for(seed): + return SimpleNamespace(weight=(1 + tensor(1, seed, 0)[0] * 0.1), variance_epsilon=EPS) + + def shared_plain(x, residual, norm): + return shared_ar.allreduce( + x, + residual=residual, + norm_weight=norm.weight, + scale=None, + bias=None, + op=shared_ar.RESIDUAL_RMS_NORM, + eps=EPS, + trigger_completion_at_end=True, + **args, + ) + + def close_pair(actual, expected, exact_residual=True): + torch.testing.assert_close( + actual[1], expected[1], rtol=0, atol=0 if exact_residual else 0.03125 + ) + # Both normalize BF16 residuals. A different CTA reduction tree may move the norm by one BF16 step. + torch.testing.assert_close(actual[0], expected[0], rtol=0.008, atol=0.0002) + + def check_state(rows, previous): + torch.cuda.synchronize() + assert state.flags.cpu().tolist() == [0, 0, (previous + 1) % 3] + assert state.layout.cpu().tolist() == [rows * HIDDEN * world, state.buffer_bytes] + + # Interleave a large shared oneshot dirty range with small owned kernels, without resets. + rows_to_test = [cell.spec["rows"] for cell in intra_node_reduce.CELLS] + [1, 18, 3] + for iteration, rows in enumerate(rows_to_test): + norm = norm_for(101 + iteration) + shared_plain(tensor(256, 11 + iteration), tensor(256, 29, 0), norm) + assert state.layout[0].item() == 256 * HIDDEN * world + x, residual = tensor(rows, 31 + iteration), tensor(rows, 73, 0) + saved_x, saved_residual = x.clone(), residual.clone() + old_flag = state.flags[2].item() + if rank == iteration % world: + time.sleep(0.02) + with validating(intra_node_reduce): + actual = ipc.plain_reduce( + norm, x, residual, trigger_completion_at_end=iteration % 2 == 0 + ) + check_state(rows, old_flag) + expected = shared_plain(x, residual, norm) + close_pair(actual, expected) + torch.testing.assert_close(x, saved_x, rtol=0, atol=0) + torch.testing.assert_close(residual, saved_residual, rtol=0, atol=0) + # Independent rank-distinct arithmetic catches a dropped or repeated rank. + reference = intra_node_reduce.reference( + [tensor(rows, 31 + iteration, r) for r in range(world)], + residual, + norm.weight, + EPS, + fp32_acc=os.getenv("ALL_REDUCE_FUSION_KERNEL_ACC_FP32", "1") != "0", + ) + close_pair(actual, reference) + total = tensor(rows, 31 + iteration, 0).float() + for other in range(1, world): + if os.getenv("ALL_REDUCE_FUSION_KERNEL_ACC_FP32", "1") != "0": + total = total + tensor(rows, 31 + iteration, other).float() + else: + total = (total.to(torch.bfloat16) + tensor(rows, 31 + iteration, other)).float() + reduced = total.to(torch.bfloat16) + residual + torch.testing.assert_close(actual[1], reduced, rtol=0, atol=0) + + def moe_payload(rows, seed): + fc2 = tensor(rows * 8, seed) + weights = torch.full((rows, 8), 0.125, dtype=torch.bfloat16, device="cuda") + indices = torch.arange(rows * 8, dtype=torch.int32, device="cuda").reshape(rows, 8) + indices[:, 3] = -1 + return fc2, weights, indices, tensor(rows, seed + 1), tensor(rows, seed + 2, 0) + + for cell in intra_node_moe_finalize.CELLS: + rows = cell.spec["rows"] + norm = norm_for(211) + fc2, weights, indices, shared, residual = moe_payload(rows, 181) + with validating(intra_node_moe_finalize): + actual = ipc.finalize(fc2, weights, indices, shared, residual, norm) + expected = torch.ops.trtllm.moe_finalize_allreduce( + input=fc2, + residual=residual, + norm_weight=norm.weight, + expanded_idx_to_permuted_idx=indices, + shared_expert_output=shared, + expert_scale_factor=weights, + workspace=table, + rank=rank, + nranks=world, + eps=EPS, + ) + close_pair(actual, expected) + reference = intra_node_moe_finalize.reference( + [tensor(rows * 8, 181, r) for r in range(world)], + [weights] * world, + [indices] * world, + [tensor(rows, 182, r) for r in range(world)], + residual, + norm.weight, + EPS, + ) + close_pair(actual, reference) + local = [] + for other in range(world): + fc = tensor(rows * 8, 181, other).reshape(rows, 8, HIDDEN) + acc = torch.zeros((rows, HIDDEN), device="cuda") + for slot in range(8): + if slot != 3: + acc = acc + fc[:, slot].float() * 0.125 + local.append((acc + tensor(rows, 182, other).float()).to(torch.bfloat16)) + reduced = local[0] + for other in range(1, world): + reduced = reduced + local[other] + torch.testing.assert_close(actual[1], reduced + residual, rtol=0, atol=0) + + scale = torch.tensor([1.0], device="cuda", dtype=torch.float32) + for cell in pre_moe_fp4_linear.CELLS: + rows = cell.spec["rows"] + norm = norm_for(303) + x, residual = tensor(rows, 281), tensor(rows, 282, 0) + with validating(pre_moe_fp4_linear): + actual = ipc.premoe(x, residual, norm, scale) + reference = pre_moe_fp4_linear.reference( + [tensor(rows, 281, r) for r in range(world)], + residual, + norm.weight, + EPS, + scale, + fp32_acc=os.getenv("ALL_REDUCE_FUSION_KERNEL_ACC_FP32", "1") != "0", + ) + close_pair(actual[:2], reference[:2]) + close_pair(actual[:2], shared_plain(x, residual, norm)) + packed, sf = torch.ops.trtllm.fp4_quantize(actual[0], scale, 16, False, False) + torch.testing.assert_close(actual[2], packed.reshape_as(actual[2]), rtol=0, atol=0) + torch.testing.assert_close( + actual[3], sf.view(torch.uint8).reshape_as(actual[3]), rtol=0, atol=0 + ) + + # Replay the same addresses through all three collective shapes, allowing flags to rotate. + for rows in (1, 6, 18): + norm = norm_for(411) + x, residual = tensor(rows, 401), tensor(rows, 402, 0) + fc2, weights, indices, shared, _ = moe_payload(rows, 403) + + def sequence(): + a = ipc.plain_reduce(norm, x, residual) + b = ipc.finalize(fc2, weights, indices, shared, a[1], norm) + # The shared FP4 kernel preloads residual before its PDL dependency wait. + # Complete this synthetic direct dependency before allowing that preload. + return ipc.premoe(b[0], b[1].clone(), norm, scale) + + sequence() + torch.cuda.synchronize() + dist.barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = sequence() + for step in range(8): + x.copy_(tensor(rows, 501 + step)) + if rank == step % world: + time.sleep(0.01) + graph.replay() + got = tuple(t.clone() for t in captured) + eager = sequence() + for a, b in zip(got, eager): + torch.testing.assert_close(a, b, rtol=0, atol=0) + assert state.flags[0].item() == 0 + del graph + + # Functionalization must fail before a peer-buffer write, not clone only the pointer table. + try: + torch.func.functionalize( + lambda a: torch.ops.trtllm.glm_intra_node_reduce( + a, residual, norm.weight, EPS, *state.args(), True + ) + )(x) + except RuntimeError as exc: + assert "functionalization" in str(exc).lower() + else: + raise AssertionError("opaque IPC functionalization unexpectedly succeeded") + dist.barrier() + torch.cuda.synchronize() + print(f"[rank {rank}] IPC arithmetic/interleave/FP4/graph checks passed", flush=True) + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker per visible device, rendezvous over env://, all killed on + the first nonzero exit or at the deadline.""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the devices this run owns" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size >= 2, f"a collective test needs at least 2 devices, got {world_size}" + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict( + os.environ, + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1", + ) + # See the allgather matrix: inheriting the launching srun step's PMI + # namespace lets one failing rank take the whole step down unreported. + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + print(f"[launcher] spawning {world_size} ranks, rendezvous on :{port}", flush=True) + children = [ + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r)), + start_new_session=True, + ) + for r in range(world_size) + ] + + def _kill_all() -> None: + for p in children: + try: + os.killpg(p.pid, signal.SIGKILL) + except ProcessLookupError: + pass + for p in children: + p.wait() + + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" + break + for r in sorted(live): + code = children[r].poll() + if code is None: + continue + live.discard(r) + if code != 0: + failure = f"rank {r} exited {code}" + break + time.sleep(0.5) + if failure is not None: + _kill_all() + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + + +if __name__ == "__main__": + if _WORKER_FLAG in sys.argv: + try: + _run_one_rank() + except BaseException: + import traceback + + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + os._exit(13) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_allgather_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/_mnnvl_allgather_op_matrix.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_allgather_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/_mnnvl_allgather_op_matrix.py index 754b92766b87..2f849093802f 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_allgather_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/_mnnvl_allgather_op_matrix.py @@ -221,7 +221,7 @@ def check_the_guard_refuses_flags_of_another_workspace() -> None: def _run_one_rank() -> None: global entry, ws, validating, RANK, WORLD, MAPPING, WORKSPACE RANK, WORLD, MAPPING = ranks.init_rank() - from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.comm import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.comm import ( # noqa: E501 _mnnvl_workspace, mnnvl_allgather, ) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_hop_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/_mnnvl_hop_op_matrix.py similarity index 76% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_hop_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/_mnnvl_hop_op_matrix.py index f8847b2712f7..91c69273f9d3 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_hop_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/_mnnvl_hop_op_matrix.py @@ -133,10 +133,18 @@ def check_certified_cells() -> None: """Every cell: the band against the reference, the same bits on every rank, and the bits of the built-in one-shot all-reduce where it runs one-shot.""" gamma = _gamma(7) - for i, cell in enumerate(entry.mnnvl_hop.CELLS): - spec = cell.spec + cells = [(cell.spec, cell.why) for cell in entry.mnnvl_hop.CELLS] + cells.extend( + ( + dict(rows=rows, hidden=HIDDEN, shard_published=published, zero_residual=zero), + f"TP8 rows={rows}, published={published}, zero_residual={zero}", + ) + for rows in (2, 4) + for published, zero in ((True, False), (False, False), (False, True)) + ) + for i, (spec, why) in enumerate(cells): rows, published, zero = spec["rows"], spec["shard_published"], spec["zero_residual"] - assert spec["hidden"] == HIDDEN, cell.why + assert spec["hidden"] == HIDDEN, why seed = 100 + 13 * i per_rank = [_rows(rows, seed, r) for r in range(WORLD)] x, residual = per_rank[RANK], _residual(rows, seed, zero) @@ -148,15 +156,15 @@ def check_certified_cells() -> None: x, residual, gamma, EPS, comm_buffer, buffer_flags, published ) ref_norm, ref_residual = entry.mnnvl_hop.reference(per_rank, residual, gamma, EPS) - _compare(norm, residual_out, ref_norm, ref_residual, cell.why) - ranks.assert_same_on_every_rank(norm, f"{cell.why}: norm_out") - ranks.assert_same_on_every_rank(residual_out, f"{cell.why}: residual_out") + _compare(norm, residual_out, ref_norm, ref_residual, why) + ranks.assert_same_on_every_rank(norm, f"{why}: norm_out") + ranks.assert_same_on_every_rank(residual_out, f"{why}: residual_out") if _builtin_one_shot(rows): builtin_norm, builtin_residual = _builtin_hop( x, residual, gamma, comm_buffer, buffer_flags ) - ranks.assert_bitwise(norm, builtin_norm, f"{cell.why}: norm_out vs built-in") - ranks.assert_bitwise(residual_out, builtin_residual, f"{cell.why}: residual_out") + ranks.assert_bitwise(norm, builtin_norm, f"{why}: norm_out vs built-in") + ranks.assert_bitwise(residual_out, builtin_residual, f"{why}: residual_out") dist.barrier() @@ -182,8 +190,7 @@ def check_exact_sums_are_exact() -> None: def check_lamport_flags_rotate() -> None: """After each hop the current buffer is the dirty one, with the hop's bytes to clear.""" gamma = _gamma(13) - for step in range(4): - rows = (8, 1, 5, 8)[step] + for step, rows in enumerate((8, 2, 4, 1, 8)): published = step % 2 == 1 _, buffer_flags = ws.hop_workspace(MAPPING, rows, HIDDEN) before = ranks.lamport_flags(buffer_flags) @@ -241,65 +248,106 @@ def check_interleaved_with_the_builtin_all_reduce() -> None: def check_cuda_graph_replays() -> None: - """A captured decode-step chain replays right: publish, published hop, hop, built-in hop. + """Replay the production o_proj publish and current-workspace hops at full/partial batch.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.gemm.o_proj_publish import ( # noqa: E501 + o_proj_publish, + ) - The flags rotate on the device, so every replay runs on the buffers the previous kernel left - current; the published rows land where the flags say at replay time. - """ gamma = _gamma(23) - rows = 8 - x_static = _rows(rows, 6000, RANK).clone() - residual_static = _residual(rows, 6000).clone() - x2_static = _rows(1, 6001, RANK).clone() - offsets = _word_offsets(rows) - - def chain(): + k = 16384 // WORLD + shapes = ((8, 2), (4, 1), (8, 1)) if WORLD == 8 else ((8, 1),) + addresses = ranks.workspace_addresses(MAPPING, HIDDEN) + states = [] + + def chain(state): + rows, draft_rows = state["shape"] + lamport, flags = ws.publish_target(MAPPING, rows, HIDDEN) + o_proj_publish(state["x"], state["weight"], state["layout"], lamport, flags, WORLD, RANK) comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, rows, HIDDEN) - _publish(x_static, offsets) norm1, res1 = entry.mnnvl_hop( - x_static, residual_static, gamma, EPS, comm_buffer, buffer_flags, True + state["layout"], state["residual"], gamma, EPS, comm_buffer, buffer_flags, True + ) + norm2, res2 = _hop(norm1, res1, gamma, False) + lamport, flags = ws.publish_target(MAPPING, draft_rows, HIDDEN) + o_proj_publish( + state["draft_x"], state["weight"], state["draft_layout"], lamport, flags, WORLD, RANK ) - norm2, res2 = entry.mnnvl_hop(norm1, res1, gamma, EPS, comm_buffer, buffer_flags, False) - comm_buffer1, buffer_flags1 = ws.hop_workspace(MAPPING, 1, HIDDEN) + comm_buffer, buffer_flags = ws.hop_workspace(MAPPING, draft_rows, HIDDEN) norm3, res3 = entry.mnnvl_hop( - x2_static, res2[:1].contiguous(), gamma, EPS, comm_buffer1, buffer_flags1, False + state["draft_layout"], + res2[:draft_rows].contiguous(), + gamma, + EPS, + comm_buffer, + buffer_flags, + True, ) return norm1, res1, norm2, res2, norm3, res3 - side = torch.cuda.Stream() - side.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(side): - chain() - torch.cuda.current_stream().wait_stream(side) - torch.cuda.synchronize() - dist.barrier() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - outputs = chain() - for replay in range(7): + for rows, draft_rows in shapes: + state = dict( + shape=(rows, draft_rows), + x=ranks.normal((rows, k), 6000 + RANK), + weight=ranks.normal((HIDDEN, k), 6010 + RANK, scale=(k * WORLD) ** -0.5), + layout=torch.full((rows, HIDDEN), 71.0, dtype=torch.bfloat16, device="cuda"), + residual=_residual(rows, 6000), + draft_x=ranks.normal((draft_rows, k), 6001 + RANK), + draft_layout=torch.full( + (draft_rows, HIDDEN), 73.0, dtype=torch.bfloat16, device="cuda" + ), + ) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + chain(state) + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + dist.barrier() + state["graph"] = torch.cuda.CUDAGraph() + with torch.cuda.graph(state["graph"]): + state["outputs"] = chain(state) + states.append(state) + assert ranks.workspace_addresses(MAPPING, HIDDEN) == addresses + + for replay in range(8): + state = states[replay % len(states)] + rows, draft_rows = state["shape"] seed = 6100 + 10 * replay - per_rank = [_rows(rows, seed, r) for r in range(WORLD)] - per_rank1 = [_rows(1, seed + 1, r) for r in range(WORLD)] - x_static.copy_(per_rank[RANK]) - residual_static.copy_(_residual(rows, seed)) - x2_static.copy_(per_rank1[RANK]) + state["x"].copy_(ranks.normal((rows, k), seed + RANK)) + state["weight"].copy_( + ranks.normal((HIDDEN, k), seed + 30 + RANK, scale=(k * WORLD) ** -0.5) + ) + state["residual"].copy_(_residual(rows, seed)) + state["draft_x"].copy_(ranks.normal((draft_rows, k), seed + 1 + RANK)) + local_published = o_proj_publish.reference( + state["x"], state["weight"], state["layout"], None, None, WORLD, RANK + ) + per_rank = [torch.empty_like(local_published) for _ in range(WORLD)] + dist.all_gather(per_rank, local_published) torch.cuda.synchronize() dist.barrier() - graph.replay() + state["graph"].replay() torch.cuda.synchronize() - norm1, res1, norm2, res2, norm3, res3 = outputs - ref = entry.mnnvl_hop.reference(per_rank, residual_static, gamma, EPS) - _compare(norm1, res1, *ref, f"replay {replay} published hop") - # The second hop sums norm1 over the ranks: build its inputs from every rank's norm1. + norm1, res1, norm2, res2, norm3, res3 = state["outputs"] + ref = entry.mnnvl_hop.reference(per_rank, state["residual"], gamma, EPS) + _compare(norm1, res1, *ref, f"replay {replay} production publish rows={rows}") + assert torch.all(state["layout"] == 71.0), "o_proj_publish wrote its layout-only output" gathered = [torch.empty_like(norm1) for _ in range(WORLD)] dist.all_gather(gathered, norm1) ref = entry.mnnvl_hop.reference(gathered, res1, gamma, EPS) _compare(norm2, res2, *ref, f"replay {replay} hop") - ref = entry.mnnvl_hop.reference(per_rank1, res2[:1], gamma, EPS) - _compare(norm3, res3, *ref, f"replay {replay} one-row hop") - ranks.assert_same_on_every_rank(norm3, f"replay {replay} one-row hop") + local_draft = o_proj_publish.reference( + state["draft_x"], state["weight"], state["draft_layout"], None, None, WORLD, RANK + ) + per_rank_draft = [torch.empty_like(local_draft) for _ in range(WORLD)] + dist.all_gather(per_rank_draft, local_draft) + assert torch.all(state["draft_layout"] == 73.0) + ref = entry.mnnvl_hop.reference(per_rank_draft, res2[:draft_rows], gamma, EPS) + _compare(norm3, res3, *ref, f"replay {replay} draft rows={draft_rows}") + ranks.assert_same_on_every_rank(norm3, f"replay {replay} draft hop") + assert ranks.workspace_addresses(MAPPING, HIDDEN) == addresses dist.barrier() - del graph + del states torch.cuda.synchronize() @@ -397,7 +445,7 @@ def check_the_guard_refuses_a_workspace_of_another_empty_value() -> None: def _run_one_rank() -> None: global entry, ws, validating, fusion_op, RANK, WORLD, MAPPING RANK, WORLD, MAPPING = ranks.init_rank() - from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.comm import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.comm import ( # noqa: E501 _mnnvl_workspace, mnnvl_hop, ) @@ -410,6 +458,7 @@ def _run_one_rank() -> None: from _validating import validating as _validating_cm validating = _validating_cm + ranks.grow_workspace_before_capture(MAPPING, HIDDEN) ranks.run_checks(CHECKS, RANK) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_moe_finalize_hop_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/_mnnvl_moe_finalize_hop_op_matrix.py similarity index 79% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_moe_finalize_hop_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/_mnnvl_moe_finalize_hop_op_matrix.py index a6c3e5478206..99633826c03c 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_moe_finalize_hop_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/_mnnvl_moe_finalize_hop_op_matrix.py @@ -138,10 +138,14 @@ def check_certified_cells() -> None: """Every cell: equal to finalize + hop bit for bit, within the band of the reference, and the same on every rank.""" gamma = _gamma(7) - for i, cell in enumerate(entry.mnnvl_moe_finalize_hop.CELLS): - spec = cell.spec + cells = [(cell.spec, cell.why) for cell in entry.mnnvl_moe_finalize_hop.CELLS] + cells.extend( + (dict(rows=rows, hidden=HIDDEN, routing=routing), f"TP8 rows={rows}, routing={routing}") + for rows, routing in ((2, "routed_row"), (4, "top8")) + ) + for i, (spec, why) in enumerate(cells): rows = spec["rows"] - assert spec["hidden"] == HIDDEN, cell.why + assert spec["hidden"] == HIDDEN, why seed = 100 + 13 * i per_rank = [_moe_inputs(rows, seed, r, spec["routing"]) for r in range(WORLD)] residual = _residual(rows, seed) @@ -162,12 +166,12 @@ def check_certified_cells() -> None: ref_norm, ref_residual = entry.mnnvl_moe_finalize_hop.reference( per_rank, residual, gamma, EPS ) - _compare(norm, residual_out, ref_norm, ref_residual, cell.why) + _compare(norm, residual_out, ref_norm, ref_residual, why) two_step_norm, two_step_residual = _two_step(per_rank[RANK], residual, gamma) - ranks.assert_bitwise(norm, two_step_norm, f"{cell.why}: norm_out vs finalize + hop") - ranks.assert_bitwise(residual_out, two_step_residual, f"{cell.why}: residual_out") - ranks.assert_same_on_every_rank(norm, f"{cell.why}: norm_out") - ranks.assert_same_on_every_rank(residual_out, f"{cell.why}: residual_out") + ranks.assert_bitwise(norm, two_step_norm, f"{why}: norm_out vs finalize + hop") + ranks.assert_bitwise(residual_out, two_step_residual, f"{why}: residual_out") + ranks.assert_same_on_every_rank(norm, f"{why}: norm_out") + ranks.assert_same_on_every_rank(residual_out, f"{why}: residual_out") dist.barrier() @@ -193,7 +197,7 @@ def check_empty_slots_and_wide_rows() -> None: def check_lamport_flags_rotate() -> None: """The fused hop rotates the workspace's flags as the plain hop does.""" gamma = _gamma(13) - for step, rows in enumerate((8, 1, 8)): + for step, rows in enumerate((8, 2, 4, 1, 8)): _, buffer_flags = ws.hop_workspace(MAPPING, rows, HIDDEN) before = ranks.lamport_flags(buffer_flags) routing = "top8" if rows > 1 else "routed_row" @@ -247,48 +251,66 @@ def check_interleaved_with_the_hops() -> None: def check_cuda_graph_replays() -> None: - """A captured verify-layer pair and a draft-step hop replay right with fresh routing.""" + """Alternate 8->2, 4->1 and 8->1 graphs with fresh routing and stable workspaces.""" gamma = _gamma(19) - static8 = [t.clone() for t in _moe_inputs(8, 5000, RANK, "top8")] - static1 = [t.clone() for t in _moe_inputs(1, 5001, RANK, "routed_row")] - residual8 = _residual(8, 5000).clone() - - def chain(): - norm8, res8 = _fused(static8, residual8, gamma) - norm1, res1 = _fused(static1, res8[:1].contiguous(), gamma) - return norm8, res8, norm1, res1 - - side = torch.cuda.Stream() - side.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(side): - chain() - torch.cuda.current_stream().wait_stream(side) - torch.cuda.synchronize() - dist.barrier() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - outputs = chain() - for replay in range(5): + shapes = ((8, 2), (4, 1), (8, 1)) if WORLD == 8 else ((8, 1),) + addresses = ranks.workspace_addresses(MAPPING, HIDDEN) + states = [] + + def chain(state): + _, draft_rows = state["shape"] + norm, residual = _fused(state["verify"], state["residual"], gamma) + draft_norm, draft_residual = _fused( + state["draft"], residual[:draft_rows].contiguous(), gamma + ) + return norm, residual, draft_norm, draft_residual + + for rows, draft_rows in shapes: + state = dict( + shape=(rows, draft_rows), + verify=[t.clone() for t in _moe_inputs(rows, 5000, RANK, "top8")], + draft=[t.clone() for t in _moe_inputs(draft_rows, 5001, RANK, "routed_row")], + residual=_residual(rows, 5000).clone(), + ) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + chain(state) + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + dist.barrier() + state["graph"] = torch.cuda.CUDAGraph() + with torch.cuda.graph(state["graph"]): + state["outputs"] = chain(state) + states.append(state) + assert ranks.workspace_addresses(MAPPING, HIDDEN) == addresses + + for replay in range(8): + state = states[replay % len(states)] + rows, draft_rows = state["shape"] seed = 5100 + 10 * replay - per_rank8 = [_moe_inputs(8, seed, r, "top8") for r in range(WORLD)] - per_rank1 = [_moe_inputs(1, seed + 1, r, "routed_row") for r in range(WORLD)] - for dst, src in zip(static8, per_rank8[RANK]): + per_rank = [_moe_inputs(rows, seed, r, "top8") for r in range(WORLD)] + per_rank_draft = [_moe_inputs(draft_rows, seed + 1, r, "routed_row") for r in range(WORLD)] + for dst, src in zip(state["verify"], per_rank[RANK]): dst.copy_(src) - for dst, src in zip(static1, per_rank1[RANK]): + for dst, src in zip(state["draft"], per_rank_draft[RANK]): dst.copy_(src) - residual8.copy_(_residual(8, seed)) + state["residual"].copy_(_residual(rows, seed)) torch.cuda.synchronize() dist.barrier() - graph.replay() + state["graph"].replay() torch.cuda.synchronize() - norm8, res8, norm1, res1 = outputs - ref = entry.mnnvl_moe_finalize_hop.reference(per_rank8, residual8, gamma, EPS) - _compare(norm8, res8, *ref, f"replay {replay} verify layer") - ref = entry.mnnvl_moe_finalize_hop.reference(per_rank1, res8[:1], gamma, EPS) - _compare(norm1, res1, *ref, f"replay {replay} draft step") - ranks.assert_same_on_every_rank(norm1, f"replay {replay} draft step") + norm, residual, draft_norm, draft_residual = state["outputs"] + ref = entry.mnnvl_moe_finalize_hop.reference(per_rank, state["residual"], gamma, EPS) + _compare(norm, residual, *ref, f"replay {replay} verify rows={rows}") + ref = entry.mnnvl_moe_finalize_hop.reference( + per_rank_draft, residual[:draft_rows], gamma, EPS + ) + _compare(draft_norm, draft_residual, *ref, f"replay {replay} draft rows={draft_rows}") + ranks.assert_same_on_every_rank(draft_norm, f"replay {replay} draft step") + assert ranks.workspace_addresses(MAPPING, HIDDEN) == addresses dist.barrier() - del graph + del states torch.cuda.synchronize() @@ -331,7 +353,7 @@ def check_the_guard_refuses_an_index_out_of_range() -> None: def _run_one_rank() -> None: global entry, hop, ws, validating, RANK, WORLD, MAPPING RANK, WORLD, MAPPING = ranks.init_rank() - from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.comm import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.comm import ( # noqa: E501 _mnnvl_workspace, mnnvl_hop, mnnvl_moe_finalize_hop, @@ -343,6 +365,7 @@ def _run_one_rank() -> None: from _validating import validating as _validating_cm validating = _validating_cm + ranks.grow_workspace_before_capture(MAPPING, HIDDEN) ranks.run_checks(CHECKS, RANK) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_ranks.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/_mnnvl_ranks.py similarity index 76% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_ranks.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/_mnnvl_ranks.py index 634311451185..81d02289ddd1 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_mnnvl_ranks.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/_mnnvl_ranks.py @@ -23,6 +23,7 @@ from __future__ import annotations +import json import os import signal import socket @@ -132,9 +133,67 @@ def init_rank() -> Tuple[int, int, Any]: gpus_per_node = min(world, torch.cuda.device_count()) mapping = Mapping(world_size=world, rank=rank, gpus_per_node=gpus_per_node, tp_size=world) + expected_world = int(os.environ.get("TLLM_TEST_EXPECT_WORLD_SIZE", world)) + assert world == expected_world, (world, expected_world) + record = dict( + rank=rank, + local_rank=local_rank, + host=socket.gethostname(), + capability=torch.cuda.get_device_capability(local_rank), + visible_gpus=torch.cuda.device_count(), + ) + records = [None] * world + dist.all_gather_object(records, record) + if os.environ.get("TLLM_TEST_EXPECT_WORLD_SIZE") == "8": + assert os.environ.get("TLLM_TEST_MNNVL", "0") != "1" + hosts = {item["host"] for item in records} + assert len(hosts) == 2, records + for host in hosts: + local = [item for item in records if item["host"] == host] + assert sorted(item["local_rank"] for item in local) == list(range(4)), local + assert all( + item["capability"] == (10, 7) and item["visible_gpus"] == 4 for item in records + ), records + print(f"[rank {rank}] topology {json.dumps(record)}", flush=True) return rank, world, mapping +def grow_workspace_before_capture(mapping, hidden: int, max_rows: int = 8) -> None: + """Prove actual manager replacement before any captured graph uses its buffers.""" + from tensorrt_llm._torch.distributed.ops import ( + MNNVLAllReduce, + get_or_scale_allreduce_mnnvl_workspace, + ) + + assert mapping not in MNNVLAllReduce.allreduce_mnnvl_workspaces, ( + "Workspace growth must run first in a fresh rank process" + ) + row_bytes = hidden * mapping.tp_size * torch.bfloat16.itemsize + small = get_or_scale_allreduce_mnnvl_workspace(mapping, torch.bfloat16, row_bytes) + assert small["buffer_size_bytes"] == row_bytes and not small["captured"] + grown = get_or_scale_allreduce_mnnvl_workspace(mapping, torch.bfloat16, max_rows * row_bytes) + assert grown is not small and grown["buffer_size_bytes"] >= max_rows * row_bytes + assert grown["uc_buffer"].data_ptr() != small["uc_buffer"].data_ptr() + assert grown["buffer_flags"].data_ptr() != small["buffer_flags"].data_ptr() + assert not grown["captured"] + print( + f"[rank {mapping.tp_rank}] workspace grew " + f"{small['buffer_size_bytes']} -> {grown['buffer_size_bytes']} bytes", + flush=True, + ) + dist.barrier() + + +def workspace_addresses(mapping, hidden: int, rows: int = 8) -> tuple[int, int]: + """Fetch the current production workspace rather than retaining an obsolete view.""" + from tensorrt_llm._torch.distributed.ops import get_or_scale_allreduce_mnnvl_workspace + + workspace = get_or_scale_allreduce_mnnvl_workspace( + mapping, torch.bfloat16, rows * hidden * mapping.tp_size * torch.bfloat16.itemsize + ) + return workspace["uc_buffer"].data_ptr(), workspace["buffer_flags"].data_ptr() + + def run_checks(checks, rank: int) -> None: """Run every check in order; a failing rank exits hard, and the launcher kills its peers.""" for check in checks: diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_rank_job.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/_rank_job.py similarity index 100% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/_rank_job.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/_rank_job.py diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/test_modeling_v2_intra_node_reduce_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/test_modeling_v2_intra_node_reduce_op_matrix.py new file mode 100644 index 000000000000..57efea4a4dde --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/test_modeling_v2_intra_node_reduce_op_matrix.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Real four-rank IPC comparison with shared kernels, dirty ranges and graph replay.""" + +import _rank_job +import pytest + +from tensorrt_llm._utils import get_sm_version + +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="owned GLM IPC kernels require SM107" +) + + +@pytest.mark.no_xdist +@pytest.mark.parametrize("cta1", ["0", "1"]) +@pytest.mark.parametrize("fp32", ["0", "1"]) +def test_intra_node_reduce_op_matrix(monkeypatch, cta1, fp32): + monkeypatch.setenv("TRTLLM_MOE_FINALIZE_CTA1", cta1) + monkeypatch.setenv("ALL_REDUCE_FUSION_KERNEL_ACC_FP32", fp32) + _rank_job.run("intra_node_reduce") diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_allgather_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/test_modeling_v2_mnnvl_allgather_op_matrix.py similarity index 100% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_allgather_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/test_modeling_v2_mnnvl_allgather_op_matrix.py diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_hop_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/test_modeling_v2_mnnvl_hop_op_matrix.py similarity index 100% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_hop_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/test_modeling_v2_mnnvl_hop_op_matrix.py diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_moe_finalize_hop_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/test_modeling_v2_mnnvl_moe_finalize_hop_op_matrix.py similarity index 100% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm/test_modeling_v2_mnnvl_moe_finalize_hop_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm/test_modeling_v2_mnnvl_moe_finalize_hop_op_matrix.py diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_glm_rank_job.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/_glm_rank_job.py similarity index 100% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_glm_rank_job.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/_glm_rank_job.py diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_hist_topk_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/_hist_topk_op_matrix.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_hist_topk_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/_hist_topk_op_matrix.py index 328e916d0087..a148bf96a351 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_hist_topk_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/_hist_topk_op_matrix.py @@ -246,10 +246,10 @@ def _run_one_rank() -> int: device_id=torch.device("cuda", local_rank), ) - from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.dsa import ( # noqa: E501 _hist_topk, ) - from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.dsa import ( # noqa: E501 hist_topk as entry, ) from tensorrt_llm.mapping import Mapping diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_indexer_inputs.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/_indexer_inputs.py similarity index 100% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/_indexer_inputs.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/_indexer_inputs.py diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_convert_req_index_to_global.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_convert_req_index_to_global.py similarity index 97% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_convert_req_index_to_global.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_convert_req_index_to_global.py index e81383ff5d33..02e380940956 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_convert_req_index_to_global.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_convert_req_index_to_global.py @@ -12,7 +12,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.convert_req_index_to_global import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.dsa.convert_req_index_to_global import ( # noqa: E501 convert_req_index_to_global as op, ) from tensorrt_llm._utils import get_sm_version diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_fp8_paged_mqa_logits.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_fp8_paged_mqa_logits.py similarity index 97% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_fp8_paged_mqa_logits.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_fp8_paged_mqa_logits.py index 56757d1a6684..22315f852a9e 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_fp8_paged_mqa_logits.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_fp8_paged_mqa_logits.py @@ -17,10 +17,10 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.fp8_paged_mqa_logits import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.dsa.fp8_paged_mqa_logits import ( # noqa: E501 fp8_paged_mqa_logits as op, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.paged_mqa_logits_schedule import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.dsa.paged_mqa_logits_schedule import ( # noqa: E501 paged_mqa_logits_schedule, ) from tensorrt_llm._utils import get_sm_version diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_full_context_topk.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_full_context_topk.py new file mode 100644 index 000000000000..160667ae7cd5 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_full_context_topk.py @@ -0,0 +1,163 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Causal full-context TopK, padded row pitches, and persistent graph workspace.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency._full_context import ( # noqa: E501 + FullContextIndexer, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.dsa._full_context_topk import ( # noqa: E501 + full_context_topk, + streaming_plan, + workspace_bytes, +) +from tensorrt_llm._torch.kernels.blackwell.top_k.gvr_topk_decode_self_sampling_host import ( + route_streaming, +) +from tensorrt_llm._utils import get_sm_version + +pytestmark = pytest.mark.skipif(get_sm_version() != 107, reason="Rubin GLM full-context indexer") + + +def test_tp8_full_context_eligibility_uses_query_window_and_live_metadata() -> None: + """TP8 admits B1/Q8 but rejects B2/Q8 and unsupported query windows.""" + capacity = 8192 + decode = SimpleNamespace(mapping=SimpleNamespace(tp_size=8)) + indexer = FullContextIndexer(decode, capacity, torch.device("cuda")) + manager = SimpleNamespace(tokens_per_block=64, use_fp4=False, max_seq_len=capacity) + for batch, next_n, supported in ( + (1, 1, True), + (2, 1, True), + (1, 4, True), + (2, 4, True), + (2, 8, False), + (1, 2, False), + (1, 6, False), + (1, 9, False), + (1, 8, True), + ): + rows = batch * next_n + metadata = SimpleNamespace( + kv_cache_manager=manager, + num_seqs=batch, + num_tokens=rows, + is_ragged_verify=False, + use_expanded_buffers_for_mtp=False, + expand_for_dsl=False, + gen_indexer_kv_lens_cuda_runtime=torch.full( + (batch,), 5000, dtype=torch.int32, device="cuda" + ), + indexer_k_cache_block_offsets=torch.zeros( + (batch, capacity // 64), dtype=torch.int32, device="cuda" + ), + shared_topk_indices=torch.empty((rows, 2048), dtype=torch.int32, device="cuda"), + scheduler_metadata_buffer=torch.empty(1, dtype=torch.int32, device="cuda"), + get_indexer_max_seq_len=lambda: manager.max_seq_len, + ) + assert indexer.supports(metadata, rows) is supported, (batch, next_n) + + # An eligible window must still have complete live metadata and fit the allocated arena. + metadata.num_tokens = 7 + assert not indexer.supports(metadata, 8) + metadata.num_tokens = 8 + metadata.shared_topk_indices = metadata.shared_topk_indices[:7] + assert not indexer.supports(metadata, 8) + metadata.shared_topk_indices = indexer.topk[:8] + manager.max_seq_len = capacity + 1 + assert not indexer.supports(metadata, 8) + + +def _assert_selection(logits, lens, got, next_n): + for row in range(logits.shape[0]): + length = min(int(lens[row // next_n]) - next_n + row % next_n + 1, logits.shape[1]) + selected = got[row] + if length <= 2048: + assert torch.equal( + selected[:length], torch.arange(length, device=got.device, dtype=torch.int32) + ) + assert (selected[length:] == -1).all() + continue + assert ((selected >= 0) & (selected < length)).all() + assert selected.unique().numel() == 2048 + expected_values = logits[row, :length].topk(2048).values + actual_values = logits[row, selected.long()].sort(descending=True).values + assert torch.equal(actual_values, expected_values) + + +@pytest.mark.parametrize( + "rows,next_n,capacity", + [ + (1, 1, 1048590), + (3, 1, 65550), + (6, 6, 65550), + (8, 8, 1048590), + (12, 6, 65550), + (18, 6, 1048590), + ], +) +@pytest.mark.parametrize("deterministic", [False, True]) +def test_causal_selection_graph_and_workspace_reuse(rows, next_n, capacity, deterministic): + generator = torch.Generator(device="cuda").manual_seed(901 + rows) + pitch = (capacity + 255) // 256 * 256 + arena = torch.randn(rows, pitch, device="cuda", generator=generator) + logits = arena[:, :capacity] + arena[:, capacity:].fill_(1e20) + batch = rows // next_n + lens = torch.tensor( + [max(next_n, 3000 + 137 * i) for i in range(batch)], device="cuda", dtype=torch.int32 + ) + indices = torch.empty(rows, 2048, device="cuda", dtype=torch.int32) + workspace = torch.zeros(workspace_bytes() // 4, device="cuda", dtype=torch.int32) + + def run(): + return full_context_topk( + logits, + lens, + indices, + workspace, + next_n=next_n, + max_seq_len=capacity, + sm_reserve=72, + deterministic=deterministic, + ) + + run() + torch.cuda.synchronize() + _assert_selection(logits, lens, indices, next_n) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + for length in (next_n, 2048, min(capacity, 33000), capacity): + lens.fill_(length) + # Many ties: every selected value must still equal the independent sorted top values. + arena[:, :capacity].copy_( + torch.randint(0, 4096, (rows, capacity), device="cuda", generator=generator).float() + ) + graph.replay() + torch.cuda.synchronize() + _assert_selection(logits, lens, indices, next_n) + if deterministic: + previous = indices.clone() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(previous, indices) + + +def test_streaming_plan_preserves_shared_abi_and_caps_all_derived_fields(): + for rows in (1, 3, 6, 8, 12, 18, 24): + for n in (129, 2048, 3000, 65550, 1048590): + tpl, splits, scap, cmp, _ = streaming_plan(rows, n, 184, 0) + baseline = route_streaming( + rows, max(n, 2049), (n + 255) // 256 * 256, 2048, force_main=True + ) + assert tpl == baseline["tpl"] + assert splits == baseline["rt"]["R"] + assert scap == baseline["rt"]["SCAP_"] and cmp == baseline["rt"]["CMP_"] + for reserve in (36, 72): + tpl, splits, _, _, _ = streaming_plan(rows, n, 184, reserve) + assert rows * splits <= 184 - reserve + assert tpl[5] == (splits > 1) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_buffer_lifecycle.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_hist_buffer_lifecycle.py similarity index 94% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_buffer_lifecycle.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_hist_buffer_lifecycle.py index 3f32aee7d0e4..0b22075ff31c 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_buffer_lifecycle.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_hist_buffer_lifecycle.py @@ -11,17 +11,17 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.modeling import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.modeling import ( # noqa: E501 DecodeTarget, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa._hist_topk import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.dsa._hist_topk import ( # noqa: E501 MAX_SHARD, shard_capacity, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.hist_shard_meta import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.dsa.hist_shard_meta import ( # noqa: E501 hist_shard_meta, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.paged_mqa_logits_schedule import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.dsa.paged_mqa_logits_schedule import ( # noqa: E501 paged_mqa_logits_schedule, ) from tensorrt_llm._utils import get_sm_version diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_shard_meta.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_hist_shard_meta.py similarity index 95% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_shard_meta.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_hist_shard_meta.py index ff3d0961e9ec..5d8d44205b12 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_shard_meta.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_hist_shard_meta.py @@ -12,10 +12,10 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa._hist_topk import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.dsa._hist_topk import ( # noqa: E501 shard_size, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.hist_shard_meta import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.dsa.hist_shard_meta import ( # noqa: E501 hist_shard_meta as op, ) from tensorrt_llm._utils import get_sm_version diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_topk_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_hist_topk_op_matrix.py similarity index 86% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_topk_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_hist_topk_op_matrix.py index d3ac1acda903..0be937dd1d9f 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_topk_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_hist_topk_op_matrix.py @@ -14,7 +14,8 @@ from tensorrt_llm._utils import get_sm_version pytestmark = pytest.mark.skipif( - get_sm_version() != 107, reason="the GLM-5.2 tp16 target's entries are certified on sm_107" + get_sm_version() != 107 or torch.cuda.device_count() < _glm_rank_job.WORLD_SIZE, + reason="requires SM107 and at least four visible GPUs", ) assert torch.cuda.is_available(), "hist_topk requires CUDA devices" diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_indexer_k_tail.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_indexer_k_tail.py similarity index 98% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_indexer_k_tail.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_indexer_k_tail.py index 534f616dc582..7e880914509c 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_indexer_k_tail.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_indexer_k_tail.py @@ -14,10 +14,10 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.dsa import ( # noqa: E501 _indexer_reference, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.indexer_k_tail import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.dsa.indexer_k_tail import ( # noqa: E501 indexer_k_tail as op, ) from tensorrt_llm._utils import get_sm_version diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_indexer_prologue.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_indexer_prologue.py similarity index 96% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_indexer_prologue.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_indexer_prologue.py index e702c25a0f32..12b5c72a248e 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_indexer_prologue.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_indexer_prologue.py @@ -13,10 +13,10 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.dsa import ( # noqa: E501 _indexer_reference, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.indexer_prologue import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.dsa.indexer_prologue import ( # noqa: E501 indexer_prologue as op, ) from tensorrt_llm._utils import get_sm_version @@ -116,7 +116,7 @@ def test_certified_cells(cell, certifying): # noqa: F811 def test_q_part_matches_the_k_tail_launch_of_the_same_kernel(): """The K part does not depend on whether the Q part is compiled in: the whole prologue and the K tail write the same K bytes for the same inputs.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.indexer_k_tail import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.dsa.indexer_k_tail import ( # noqa: E501 indexer_k_tail, ) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_paged_mqa_logits_schedule.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_paged_mqa_logits_schedule.py similarity index 96% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_paged_mqa_logits_schedule.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_paged_mqa_logits_schedule.py index cbbc5b6b61ba..3b001bd907f7 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_paged_mqa_logits_schedule.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_paged_mqa_logits_schedule.py @@ -8,7 +8,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.paged_mqa_logits_schedule import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.dsa.paged_mqa_logits_schedule import ( # noqa: E501 paged_mqa_logits_schedule as op, ) from tensorrt_llm._utils import get_sm_version diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_stash_last_accepted_rows.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_stash_last_accepted_rows.py similarity index 62% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_stash_last_accepted_rows.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_stash_last_accepted_rows.py index c8593238d1bd..fae35bf6f658 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_stash_last_accepted_rows.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/dsa/test_modeling_v2_stash_last_accepted_rows.py @@ -11,7 +11,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.dsa.stash_last_accepted_rows import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.dsa.stash_last_accepted_rows import ( # noqa: E501 stash_last_accepted_rows as op, ) from tensorrt_llm._utils import get_sm_version @@ -54,31 +54,51 @@ def test_certified_cells(cell, certifying): # noqa: F811 assert (shared[cell.spec["requests"] :] == -7).all(), "rows past the requests moved" -def test_rows_follow_the_accepted_counts_in_place_under_replay(): +@pytest.mark.parametrize( + "requests,next_n,replay_counts", + [ + (1, 8, [(1,), (6,), (8,), (0,), (9,)]), + (3, 6, [(1, 4, 6), (6, 1, 4), (0, 7, -2), (8, 2, 3)]), + (2, 4, [(1, 4), (4, 1), (2, 3), (0, 5)]), + (1, 4, [(1,), (4,), (2,), (0,), (5,)]), + ], +) +def test_rows_follow_the_accepted_counts_in_place_under_replay(requests, next_n, replay_counts): """Captured once, replayed with new counts: the counts are read on the device.""" - args, _ = _build(op.CELLS[0].spec, seed=71) + spec = dict(requests=requests, next_n=next_n, n_cols=2048, accepted=(1,) * requests, dst="wide") + args, shared = _build(spec, seed=71) op(**args) torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): op(**args) - for accepted in (1, 6, 8): - args["num_accepted"].fill_(accepted) + for accepted in replay_counts: + args["num_accepted"].copy_(torch.tensor(accepted, dtype=torch.int32, device="cuda")) graph.replay() torch.cuda.synchronize() - assert torch.equal(args["dst"][0], args["gen_topk"][accepted - 1]) + # Select within each request's own block, independently of the op's reference. + expected = torch.stack( + [ + rows[min(max(count, 1), next_n) - 1] + for rows, count in zip(args["gen_topk"].split(next_n), accepted, strict=True) + ] + ) + assert torch.equal(args["dst"], expected) + assert (shared[requests:] == -7).all(), "rows past the requests moved" def test_guard_refuses_a_topk_of_another_row_count(): - args, _ = _build(op.CELLS[0].spec, seed=72) + spec = dict(requests=3, next_n=6, n_cols=2048, accepted=(1, 4, 6), dst="rows") + args, _ = _build(spec, seed=72) with pytest.raises(AssertionError, match="gen_topk must be"): with validating(op): - op(**dict(args, gen_topk=args["gen_topk"][:4])) + op(**dict(args, gen_topk=args["gen_topk"][:-1])) def test_guard_refuses_a_strided_destination(): - args, _ = _build(op.CELLS[0].spec, seed=73) - wide = torch.zeros(1, 2 * 2048, dtype=torch.int32, device="cuda") + spec = dict(requests=3, next_n=6, n_cols=2048, accepted=(1, 4, 6), dst="rows") + args, _ = _build(spec, seed=73) + wide = torch.zeros(spec["requests"], 2 * spec["n_cols"], dtype=torch.int32, device="cuda") with pytest.raises(AssertionError, match="unit column strides"): with validating(op): op(**dict(args, dst=wide[:, ::2])) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/embedding/test_modeling_v2_fused_embedding.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/embedding/test_modeling_v2_fused_embedding.py similarity index 97% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/embedding/test_modeling_v2_fused_embedding.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/embedding/test_modeling_v2_fused_embedding.py index 555e0509d8dc..8499db1671ba 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/embedding/test_modeling_v2_fused_embedding.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/embedding/test_modeling_v2_fused_embedding.py @@ -13,11 +13,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.embedding.fused_embedding import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.embedding.fused_embedding import ( # noqa: E501 HIDDEN, VOCAB, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.embedding.fused_embedding import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.embedding.fused_embedding import ( # noqa: E501 fused_embedding as op, ) from tensorrt_llm._torch.modules.embedding import pre_comm_embedding_ops diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_dense_mlp_small_m.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_dense_mlp_small_m.py similarity index 98% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_dense_mlp_small_m.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_dense_mlp_small_m.py index 98184625ba61..10adafda9107 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_dense_mlp_small_m.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_dense_mlp_small_m.py @@ -14,11 +14,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.dense_mlp_small_m import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.gemm.dense_mlp_small_m import ( # noqa: E501 HIDDEN, INTER_PER_RANK, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.dense_mlp_small_m import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.gemm.dense_mlp_small_m import ( # noqa: E501 dense_mlp_small_m as op, ) from tensorrt_llm._utils import get_sm_version diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_draft_hidden_half.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_draft_hidden_half.py similarity index 69% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_draft_hidden_half.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_draft_hidden_half.py index 94966647d7c4..01f069b5b161 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_draft_hidden_half.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_draft_hidden_half.py @@ -4,12 +4,16 @@ import os import zlib +from types import SimpleNamespace import pytest import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.draft_hidden_half import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency._small_batch_draft import ( # noqa: E501 + SmallBatchDraftKernels, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.gemm.draft_hidden_half import ( # noqa: E501 draft_hidden_half as op, ) from tensorrt_llm._utils import get_sm_version @@ -38,7 +42,9 @@ def _build(spec: dict, seed: int) -> dict: if not in_kernel_norm: # The fallback's input: the row(s) already normalized and rounded to BF16. xf = x.float() - x = (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + spec["eps"])).to(torch.bfloat16) + x = ( + xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + spec["eps"]) * norm_weight.float() + ).to(torch.bfloat16) return dict( x=x, weight=weight, @@ -93,12 +99,13 @@ def _flashinfer_norm_is_cute_dsl() -> bool: return not _use_cuda_norm(torch.device("cuda", torch.cuda.current_device())) -def test_in_kernel_norm_is_flashinfer_rmsnorm_bit_for_bit() -> None: +@pytest.mark.parametrize("n", [384, 768]) +def test_in_kernel_norm_is_flashinfer_rmsnorm_bit_for_bit(n) -> None: """The point of the in-kernel norm: the result of FlashInfer's RMSNorm followed by this GEMM, bit for bit, from one kernel.""" if not _flashinfer_norm_is_cute_dsl(): pytest.skip("FlashInfer runs its CUDA RMSNorm kernel here, not the one reproduced") - args = _build(op.CELLS[0].spec, seed=71) + args = _build(dict(op.CELLS[0].spec, n=n), seed=71) with validating(op): fused = op(**args) normed = torch.ops.trtllm.flashinfer_rmsnorm(args["x"], args["norm_weight"], args["eps"]) @@ -106,10 +113,11 @@ def test_in_kernel_norm_is_flashinfer_rmsnorm_bit_for_bit() -> None: assert torch.equal(fused, plain) -def test_cuda_graph_capture_and_replay() -> None: +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_cuda_graph_capture_and_replay(cell) -> None: """The draft loop captures the op on the auxiliary stream: it must replay on new contents of the same buffers. Warm-up compiles the kernel outside the capture.""" - spec = op.CELLS[0].spec + spec = cell.spec args = _build(spec, seed=72) warmup = torch.cuda.Stream() @@ -126,7 +134,54 @@ def test_cuda_graph_capture_and_replay() -> None: args["out"].fill_(float("nan")) graph.replay() torch.cuda.synchronize() - _compare(args["out"], op.reference(**args), in_kernel_norm=True) + _compare(args["out"], op.reference(**args), in_kernel_norm=spec["in_kernel_norm"]) + + +@pytest.mark.parametrize("rows", [1, 2]) +def test_tp8_head_auxiliary_half_join_replays_new_hidden_states(rows): + """The production head keeps weighted hnorm and its FP32 half alive until the main-stream join.""" + args = _build(dict(op.CELLS[0].spec, n=768, rows=rows), seed=91) + helper = object.__new__(SmallBatchDraftKernels) + helper.use_table = True + helper.hidden_rows = args["weight"] + helper.hnorm_weight = args["norm_weight"] + helper.half_start, helper.half_done = torch.cuda.Event(), torch.cuda.Event() + + class Norm: + variance_epsilon = args["eps"] + + def __call__(self, hidden): + return torch.ops.trtllm.flashinfer_rmsnorm( + hidden, helper.hnorm_weight, self.variance_epsilon + ) + + helper.layer = SimpleNamespace( + mlp=SimpleNamespace(aux_stream=torch.cuda.Stream()), hnorm=Norm() + ) + helper.lm_head = object() + helper.decode = SimpleNamespace(_lm_head_shard=lambda lm, hidden: hidden.square()) + + def run(): + logits, half = helper.head(args["x"], next_step=True) + torch.cuda.current_stream().wait_event(helper.half_done) + return logits, half.clone() + + run() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + logits, result = run() + for seed in (92, 93): + args["x"].copy_(_build(dict(op.CELLS[0].spec, n=768, rows=rows), seed=seed)["x"]) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(logits, args["x"].square(), rtol=0, atol=0) + if rows == 1: + reference = op.reference(**args) + else: + normed = helper.layer.hnorm(args["x"]) + reference = op.reference(normed, args["weight"], result, None, None) + _compare(result, reference, in_kernel_norm=rows == 1) def test_guard_refuses_eps_without_the_norm_weight() -> None: diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_fused_a_gemm.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_fused_a_gemm.py similarity index 95% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_fused_a_gemm.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_fused_a_gemm.py index eed24372acbe..1d1d1d9be391 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_fused_a_gemm.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_fused_a_gemm.py @@ -13,12 +13,12 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.fused_a_gemm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.gemm.fused_a_gemm import ( # noqa: E501 HIDDEN_SIZE, INDEXER_OUT_FEATURES, OUT_FEATURES, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.fused_a_gemm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.gemm.fused_a_gemm import ( # noqa: E501 fused_a_gemm as op, ) from tensorrt_llm._utils import get_sm_version @@ -174,7 +174,7 @@ def test_op_rejects_what_the_kernel_cannot_run() -> None: ("3D hidden_states", lambda: op(hidden_states.unsqueeze(0), weight, [])), ("hidden 7168", lambda: op(hidden_7168, weight_7168, [])), ("2640 rows", lambda: op(hidden_states, _tensor((2640, HIDDEN_SIZE), BF16, g), [])), - ("9 tokens", lambda: op(_tensor((9, HIDDEN_SIZE), BF16, g), weight, [])), + ("33 tokens", lambda: op(_tensor((33, HIDDEN_SIZE), BF16, g), weight, [])), ("0 tokens", lambda: op(hidden_states[:0], weight, [])), ("strided hidden_states", lambda: op(strided_hidden_states, weight, [])), ("strided weight", lambda: op(hidden_states, strided_weight, [])), @@ -186,7 +186,7 @@ def test_op_rejects_what_the_kernel_cannot_run() -> None: ("misaligned list tensor", lambda: op(hidden_states, weight, [misaligned])), ] for tag, call in cases: - with pytest.raises((RuntimeError, NotImplementedError)): + with pytest.raises((RuntimeError, NotImplementedError, ValueError)): call() pytest.fail(f"{tag} was accepted") @@ -215,11 +215,11 @@ def test_guard_refuses_a_weight_without_an_instance() -> None: op(hidden_states, _tensor((2640, HIDDEN_SIZE), BF16, g), []) -@pytest.mark.parametrize("tokens", [0, 9, 16]) -def test_guard_refuses_a_token_count_outside_one_tile(tokens) -> None: +@pytest.mark.parametrize("tokens", [0, 33]) +def test_guard_refuses_a_token_count_outside_supported_range(tokens) -> None: hidden_states, weight = _ok(tokens=max(tokens, 1)) hidden_states = hidden_states[:tokens] - with pytest.raises(AssertionError, match="1 to 8 tokens"): + with pytest.raises(AssertionError, match="1 to 32 tokens"): with validating(op): op(hidden_states, weight, []) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_indexer_q_tail.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_indexer_q_tail.py similarity index 98% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_indexer_q_tail.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_indexer_q_tail.py index 03675a813a89..38aa4b4fb812 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_indexer_q_tail.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_indexer_q_tail.py @@ -7,7 +7,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.indexer_q_tail import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.gemm.indexer_q_tail import ( # noqa: E501 indexer_q_tail as op, ) from tensorrt_llm._utils import get_sm_version diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_low_m_gemm.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_low_m_gemm.py similarity index 57% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_low_m_gemm.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_low_m_gemm.py index ffe9d1b60579..446183d6b7bc 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_low_m_gemm.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_low_m_gemm.py @@ -7,7 +7,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.low_m_gemm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.gemm.low_m_gemm import ( # noqa: E501 low_m_gemm as op, ) from tensorrt_llm._utils import get_sm_version @@ -59,21 +59,29 @@ def test_certified_cells(cell, certifying): # noqa: F811 _compare(out, op.reference(x, weight, out)) -def test_cuda_graph_capture_and_replay() -> None: +@pytest.mark.parametrize( + "spec, one_wave", + [ + (dict(m=8, k=6144, n=256), False), + (dict(m=8, k=6144, n=19360), True), + (dict(m=8, k=6144, n=3072), False), + (dict(m=8, k=1536, n=6144), False), + ], +) +def test_cuda_graph_capture_and_replay(spec, one_wave) -> None: """The target captures the whole decode step: the op must replay on new contents of the same buffers. Warm-up compiles the kernel outside the capture.""" - spec = op.CELLS[0].spec x, weight, out = _build(spec, seed=31) warmup = torch.cuda.Stream() warmup.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(warmup): - op(x, weight, out) + op(x, weight, out, one_wave=one_wave) torch.cuda.current_stream().wait_stream(warmup) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - op(x, weight, out) + op(x, weight, out, one_wave=one_wave) fresh, _, _ = _build(spec, seed=32) x.copy_(fresh) @@ -102,3 +110,70 @@ def test_guard_refuses_an_out_over_an_input(operand) -> None: with pytest.raises(AssertionError, match=f"out must not share memory with {operand}"): with validating(op): op(x, weight, out) + + +@pytest.mark.parametrize("rows", [1, 2, 3, 6, 12, 18]) +def test_batched_v_b_strides_and_graph(rows): + from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import true_fp32_matmul + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.gemm._low_m_splitk import ( # noqa: E501 + bmm_prewait_tactic, + run_splitk_bmm, + ) + + gen = torch.Generator(device="cuda").manual_seed(400 + rows) + latent = torch.randn(rows, 16, 512, device="cuda", generator=gen).bfloat16() + weight = (torch.randn(16, 256, 512, device="cuda", generator=gen) / 512**0.5).bfloat16() + storage = torch.full((rows + 2, 16, 256), float("nan"), device="cuda", dtype=torch.bfloat16) + output = storage[1:-1] + tactic = bmm_prewait_tactic(rows, 256, 512) + + def run(): + return run_splitk_bmm( + latent.transpose(0, 1), + weight, + output.transpose(0, 1), + True, + tactic, + reserve_full_smem=False, + ) + + run() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + for _ in range(3): + latent.copy_(torch.randn(latent.shape, device="cuda", generator=gen).bfloat16()) + graph.replay() + torch.cuda.synchronize() + with true_fp32_matmul(): + ref = torch.bmm(latent.transpose(0, 1).float(), weight.float().transpose(1, 2)) + _compare(output, ref.transpose(0, 1).bfloat16()) + assert storage[0].isnan().all() and storage[-1].isnan().all() + + +@pytest.mark.parametrize( + "rows, vocab", [(3, 38720), (1, 19360), (2, 19360), (4, 19360), (8, 19360)] +) +def test_one_wave_lm_head_numeric(rows, vocab): + x, weight, out = _build(dict(m=rows, k=6144, n=vocab), seed=715) + op(x, weight, out, one_wave=True) + _compare(out, op.reference(x, weight, out)) + + +@pytest.mark.parametrize( + "rows, inter", [(6, 3072), (12, 3072), (18, 3072), (1, 1536), (2, 1536), (4, 1536), (8, 1536)] +) +def test_dense_mlp_composition(rows, inter): + from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import true_fp32_matmul + from tensorrt_llm._torch.modules.swiglu import swiglu + + x, up, mid = _build(dict(m=rows, k=6144, n=2 * inter), seed=800 + rows) + _, down, out = _build(dict(m=rows, k=inter, n=6144), seed=900 + rows) + got = op(swiglu(op(x, up, mid)), down, out) + with true_fp32_matmul(): + gate_up = (x.float() @ up.float().t()).bfloat16() + gate, value = gate_up.float().chunk(2, dim=-1) + activated = (torch.nn.functional.silu(gate) * value).bfloat16() + ref = (activated.float() @ down.float().t()).bfloat16() + _compare(got, ref) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_o_proj_publish.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_o_proj_publish.py similarity index 97% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_o_proj_publish.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_o_proj_publish.py index 6a856d09fe8e..ecddb8b0e47d 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_o_proj_publish.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_o_proj_publish.py @@ -13,10 +13,10 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.low_m_gemm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.gemm.low_m_gemm import ( # noqa: E501 low_m_gemm, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.o_proj_publish import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.gemm.o_proj_publish import ( # noqa: E501 o_proj_publish as op, ) from tensorrt_llm._utils import get_sm_version diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_q_b_fp8_rope.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_q_b_fp8_rope.py similarity index 98% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_q_b_fp8_rope.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_q_b_fp8_rope.py index cd70e5d27797..56ea7c5ee05f 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_q_b_fp8_rope.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_q_b_fp8_rope.py @@ -7,7 +7,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.q_b_fp8_rope import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.gemm.q_b_fp8_rope import ( # noqa: E501 q_b_fp8_rope as op, ) from tensorrt_llm._utils import get_sm_version diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_router_gemm.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_router_gemm.py similarity index 87% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_router_gemm.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_router_gemm.py index 29a981ee05a2..da2a708b7177 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/gemm/test_modeling_v2_router_gemm.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/gemm/test_modeling_v2_router_gemm.py @@ -12,12 +12,12 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.router_gemm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.gemm.router_gemm import ( # noqa: E501 HIDDEN_SIZE, MAX_TOKENS, NUM_EXPERTS, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.gemm.router_gemm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.gemm.router_gemm import ( # noqa: E501 router_gemm as op, ) from tensorrt_llm._utils import get_sm_version @@ -80,7 +80,7 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res # ── the claims no cell carries ──────────────────────────────────────────── -@pytest.mark.parametrize("tokens", range(1, MAX_TOKENS + 1)) +@pytest.mark.parametrize("tokens", range(1, 17)) def test_bitwise_the_shared_router_gemv(tokens) -> None: """dsv3_router_gemm_op runs the same loads and FMA chain; only the trigger and the L2 prefetch differ, and neither touches the arithmetic. Every token count is an instance.""" @@ -91,23 +91,25 @@ def test_bitwise_the_shared_router_gemv(tokens) -> None: _compare(logits, op.reference(hidden_states, weight)) -def test_deterministic() -> None: - hidden_states, weight = _build(8, seed=21) +@pytest.mark.parametrize("tokens", [8, 17, 18, 24, 32]) +def test_deterministic(tokens) -> None: + hidden_states, weight = _build(tokens, seed=21) first = op(hidden_states, weight) second = op(hidden_states, weight) assert torch.equal(first, second) and first.data_ptr() != second.data_ptr() -def test_cuda_graph_replay() -> None: +@pytest.mark.parametrize("tokens", [1, 18, 32]) +def test_cuda_graph_replay(tokens) -> None: """The decode step runs as one CUDA graph: the captured activations change between replays.""" - hidden_states, weight = _build(1, seed=4) + hidden_states, weight = _build(tokens, seed=4) op(hidden_states, weight) torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): logits = op(hidden_states, weight) for seed in (5, 6): - hidden_states.copy_(_build(1, seed)[0]) + hidden_states.copy_(_build(tokens, seed)[0]) graph.replay() torch.cuda.synchronize() _compare(logits, op.reference(hidden_states, weight)) @@ -123,7 +125,7 @@ def test_op_rejects_what_the_kernel_cannot_run() -> None: hidden_7168 = torch.randn(8, 7168, generator=g, device=DEV).to(BF16) weight_7168 = torch.randn(NUM_EXPERTS, 7168, generator=g, device=DEV).to(BF16) weight_128 = weight[:128] - hidden_17 = torch.randn(17, HIDDEN_SIZE, generator=g, device=DEV).to(BF16) + hidden_33 = torch.randn(33, HIDDEN_SIZE, generator=g, device=DEV).to(BF16) strided_hidden_states = torch.randn(HIDDEN_SIZE, 8, generator=g, device=DEV).to(BF16).t() flat = torch.zeros(8 * HIDDEN_SIZE + 8, dtype=BF16, device=DEV) misaligned = flat[1 : 1 + 8 * HIDDEN_SIZE].view(8, HIDDEN_SIZE) @@ -133,7 +135,7 @@ def test_op_rejects_what_the_kernel_cannot_run() -> None: ("3D hidden_states", lambda: op(hidden_states.unsqueeze(0), weight)), ("hidden 7168", lambda: op(hidden_7168, weight_7168)), ("128 experts", lambda: op(hidden_states, weight_128)), - ("17 tokens", lambda: op(hidden_17, weight)), + ("33 tokens", lambda: op(hidden_33, weight)), ("0 tokens", lambda: op(hidden_states[:0], weight)), ("strided hidden_states", lambda: op(strided_hidden_states, weight)), ("weight.t()", lambda: op(hidden_states, weight.t())), @@ -165,9 +167,15 @@ def test_guard_refuses_another_expert_count() -> None: op(hidden_states, weight[:128]) -@pytest.mark.parametrize("tokens", [0, 17, 32]) +@pytest.mark.parametrize("tokens", [0, 33, 64]) def test_guard_refuses_a_token_count_without_an_instance(tokens) -> None: hidden_states, weight = _build(max(tokens, 1), seed=3) - with pytest.raises(AssertionError, match="1 to 16 tokens"): + with pytest.raises(AssertionError, match="1 to 32 tokens"): with validating(op): op(hidden_states[:tokens], weight) + + +@pytest.mark.parametrize("tokens", range(17, MAX_TOKENS + 1)) +def test_tensorcore_router_all_token_counts(tokens) -> None: + hidden_states, weight = _build(tokens, seed=100 + tokens) + _compare(op(hidden_states, weight), op.reference(hidden_states, weight)) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_lowm_bf16_moe.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/moe/test_modeling_v2_lowm_bf16_moe.py similarity index 79% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_lowm_bf16_moe.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/moe/test_modeling_v2_lowm_bf16_moe.py index 9c609339c90c..60f8eb0749ef 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_lowm_bf16_moe.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/moe/test_modeling_v2_lowm_bf16_moe.py @@ -7,7 +7,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.moe.lowm_bf16_moe import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.moe.lowm_bf16_moe import ( # noqa: E501 glm_lowm_bf16_moe as op, ) from tensorrt_llm._utils import get_sm_version @@ -49,7 +49,8 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: def _build(spec, seed): g = torch.Generator(device="cuda").manual_seed(seed) hidden, inter = spec["hidden"], spec["inter"] - x = torch.randn(1, hidden, generator=g, device="cuda", dtype=torch.bfloat16) + tokens = spec.get("tokens", 1) + x = torch.randn(tokens, hidden, generator=g, device="cuda", dtype=torch.bfloat16) w3_w1 = torch.randn( NUM_EXPERTS, 2 * inter, hidden, generator=g, device="cuda", dtype=torch.bfloat16 ).mul_(hidden**-0.5) @@ -58,17 +59,19 @@ def _build(spec, seed): ).mul_(inter**-0.5) routing = spec["routing"] if routing == "random": - logits = 2.0 * torch.randn(1, NUM_EXPERTS, generator=g, device="cuda") + logits = 2.0 * torch.randn(tokens, NUM_EXPERTS, generator=g, device="cuda") bias = 0.1 * torch.randn(NUM_EXPERTS, generator=g, device="cuda") elif routing == "ties": # Half-integer logits in [-2, 2] and no bias: about 28 experts share the top key, so which 8 of them are # selected, and in which order, is the tie-break. - logits = torch.randint(-4, 5, (1, NUM_EXPERTS), generator=g, device="cuda").float() / 2.0 + logits = ( + torch.randint(-4, 5, (tokens, NUM_EXPERTS), generator=g, device="cuda").float() / 2.0 + ) bias = torch.zeros(NUM_EXPERTS, device="cuda") else: assert routing == "tiny" # sigmoid(-16 +- 1) is about 1e-7, so the 8 selected scores sum to about 1e-6 < 2^-13. - logits = -16.0 + 0.5 * torch.randn(1, NUM_EXPERTS, generator=g, device="cuda") + logits = -16.0 + 0.5 * torch.randn(tokens, NUM_EXPERTS, generator=g, device="cuda") bias = 0.1 * torch.randn(NUM_EXPERTS, generator=g, device="cuda") return x, logits.contiguous(), bias.float().contiguous(), w3_w1, w2 @@ -82,9 +85,10 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by _compare(out, op.reference(*args, TOP_K, ROUTED_SCALING_FACTOR)) -def test_cuda_graph_replay_matches_eager() -> None: +@pytest.mark.parametrize("tokens,inter", [(1, 128), (2, 512), (3, 512), (1, 256), (2, 256)]) +def test_cuda_graph_replay_matches_eager(tokens, inter) -> None: """The draft loop runs inside the decode graph: a replay with new inputs gives the eager result bitwise.""" - spec = op.CELLS[0].spec + spec = dict(op.CELLS[0].spec, tokens=tokens, inter=inter) x, logits, bias, w3_w1, w2 = _build(spec, seed=1) static_x, static_logits = x.clone(), logits.clone() out = op(static_x, static_logits, bias, w3_w1, w2, TOP_K, ROUTED_SCALING_FACTOR) @@ -115,14 +119,11 @@ def _small(hidden=2048, inter=64, tokens=1): return x, logits, bias, w3_w1, w2 -def test_refuses_more_than_one_token() -> None: - """Draft step 0 without the row select has 8 rows: the caller must take another MoE path. - - The op says so rather than computing token 0 alone. - """ - with pytest.raises(RuntimeError, match="takes one token"): +def test_refuses_more_than_six_tokens() -> None: + """Verify rows need another MoE path; no silent token-zero-only result.""" + with pytest.raises(RuntimeError, match="takes 1 to 6 tokens"): with validating(op): - op(*_small(tokens=2), TOP_K, ROUTED_SCALING_FACTOR) + op(*_small(tokens=7), TOP_K, ROUTED_SCALING_FACTOR) def test_refuses_an_expert_slice_it_does_not_serve() -> None: @@ -147,3 +148,18 @@ def test_refuses_a_misaligned_token() -> None: with pytest.raises(RuntimeError, match="16-byte aligned"): with validating(op): op(misaligned, logits, bias, w3_w1, w2, TOP_K, ROUTED_SCALING_FACTOR) + + +@pytest.mark.parametrize("tokens,inter", [(2, 512), (3, 512), (2, 256)]) +def test_multi_matches_single_rows_bitwise(tokens, inter) -> None: + args = _build(dict(tokens=tokens, hidden=6144, inter=inter, routing="random"), seed=32 + tokens) + x, logits, bias, w3_w1, w2 = args + multi = op(*args, TOP_K, ROUTED_SCALING_FACTOR) + single = torch.cat( + [ + op(x[r : r + 1], logits[r : r + 1], bias, w3_w1, w2, TOP_K, ROUTED_SCALING_FACTOR) + for r in range(tokens) + ] + ) + torch.cuda.synchronize() + assert torch.equal(multi, single) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_routed_experts.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/moe/test_modeling_v2_routed_experts.py similarity index 81% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_routed_experts.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/moe/test_modeling_v2_routed_experts.py index e9774111aa57..b90c76bb653b 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/moe/test_modeling_v2_routed_experts.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/moe/test_modeling_v2_routed_experts.py @@ -9,13 +9,13 @@ from torch.profiler import ProfilerActivity, profile from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import true_fp32_matmul -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.moe.routed_experts import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.moe.routed_experts import ( # noqa: E501 ACT_TYPE_SWIGLU, FUSED_ROUTER_WORKSPACE_WORDS, ROUTING_METHOD_DEEPSEEK_V3, make_fused_router_workspace, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.moe.routed_experts import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.moe.routed_experts import ( # noqa: E501 routed_experts as op, ) from tensorrt_llm._utils import get_sm_version @@ -110,22 +110,21 @@ def _prepare(packed: torch.Tensor, scales: torch.Tensor, perm: torch.Tensor): return w, _swizzle(s).contiguous().view(torch.float8_e4m3fn) -_STACK = None +_STACK = {} -def _stack(): +def _stack(inter=INTER): """GLM-5.2's expert stack at MoE TP16 -- 256 experts, H 6144, I 128 -- built once: the kernel operands and the reference's `(codes, scales)`.""" - global _STACK - if _STACK is None: + if inter not in _STACK: gen = torch.Generator(device=DEV).manual_seed(5216) - up_p, up_c, up_s = _rand_nvfp4(EXPERTS, INTER, HIDDEN, gen) - gt_p, gt_c, gt_s = _rand_nvfp4(EXPERTS, INTER, HIDDEN, gen) - dn_p, dn_c, dn_s = _rand_nvfp4(EXPERTS, HIDDEN, INTER, gen) - fc1_perm = _interleave_perm(2 * INTER)[_blk32_perm(2 * INTER)] + up_p, up_c, up_s = _rand_nvfp4(EXPERTS, inter, HIDDEN, gen) + gt_p, gt_c, gt_s = _rand_nvfp4(EXPERTS, inter, HIDDEN, gen) + dn_p, dn_c, dn_s = _rand_nvfp4(EXPERTS, HIDDEN, inter, gen) + fc1_perm = _interleave_perm(2 * inter)[_blk32_perm(2 * inter)] w1, s1 = _prepare(torch.cat([up_p, gt_p], dim=1), torch.cat([up_s, gt_s], dim=1), fc1_perm) w2, s2 = _prepare(dn_p, dn_s, _blk32_perm(HIDDEN)) - _STACK = dict( + _STACK[inter] = dict( gemm1_weights=w1, gemm1_weights_scale=s1, gemm2_weights=w2, @@ -134,7 +133,7 @@ def _stack(): gate=(gt_c, gt_s), down=(dn_c, dn_s), ) - return _STACK + return _STACK[inter] def _router(gen: torch.Generator): @@ -166,10 +165,10 @@ def _quantize(x: torch.Tensor, g1: torch.Tensor): return xq, xsf.view(torch.float8_e4m3fn) -def _scalars(xq, xsf, chosen, g1: torch.Tensor): +def _scalars(xq, xsf, chosen, g1: torch.Tensor, inter=INTER): """The three per-expert scalars for weights with a global scale of 1, and `g2` calibrated on the true FC1 activation of the chosen experts.""" - stack = _stack() + stack = _stack(inter) x = _dequant_act(xq, xsf) / g1 amax = torch.zeros((), device=DEV) with true_fp32_matmul(): @@ -217,7 +216,7 @@ def _selection(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor): return logits, top.indices[:, :TOP_K], margin -def _case(tokens: int, seed: int, clear_margin: bool = False) -> dict: +def _case(tokens: int, seed: int, clear_margin: bool = False, inter=INTER) -> dict: """One call's inputs. `clear_margin` redraws (deterministically, from the next seeds) a router input whose expert choice sits within `SELECTION_MARGIN` of a tie, which only a comparison against the @@ -233,7 +232,7 @@ def _case(tokens: int, seed: int, clear_margin: bool = False) -> dict: raise AssertionError(f"no draw from seed {seed} clears the selection margin") g1 = _input_global_scale(x) xq, xsf = _quantize(x, g1) - s1, sg, s2 = _scalars(xq, xsf, chosen, g1) + s1, sg, s2 = _scalars(xq, xsf, chosen, g1, inter) return dict( x=x, g1=g1, @@ -245,6 +244,7 @@ def _case(tokens: int, seed: int, clear_margin: bool = False) -> dict: output1_scale_scalar=s1, output1_scale_gate_scalar=sg, output2_scale_scalar=s2, + **({"intermediate": inter} if inter != INTER else {}), ) @@ -259,13 +259,14 @@ def _call_args( fused_router: bool, workspace=None, event=None, - gemm2_k128: bool, + specialized_gemm2: bool, do_finalize: bool = False, routing_logits=None, output=None, ): """The target's call, every argument explicit.""" - stack = _stack() + inter = case.get("intermediate", INTER) + stack = _stack(inter) tokens = case["x"].shape[0] if routing_logits is None: # With the fused router only the shape is read; filled so that an @@ -294,7 +295,7 @@ def _call_args( top_k=TOP_K, n_group=1, topk_group=1, - intermediate_size=INTER, + intermediate_size=inter, local_expert_offset=0, local_num_experts=EXPERTS, routed_scaling_factor=ROUTED_SCALE, @@ -310,7 +311,7 @@ def _call_args( router_hidden_states=case["x"] if fused_router else None, router_weight=case["router_weight"] if fused_router else None, router_workspace=workspace if fused_router else None, - prefer_gemm2_k128=gemm2_k128, + prefer_specialized_gemm2=specialized_gemm2, ) @@ -354,6 +355,9 @@ def _finalize(outs) -> torch.Tensor: def _assert_same(a, b, what: str) -> None: + if len(a) == len(b) == 1: + assert torch.equal(a[0], b[0]), f"{what}: finalized outputs differ" + return for x, y, name in zip(_valid_rows(a), _valid_rows(b), ("rows", "weights", "map")): assert torch.equal(x, y), f"{what}: {name} differ" @@ -372,13 +376,14 @@ def _kernels(fn) -> list[str]: return [e.name for e in sorted(launched, key=lambda e: e.time_range.start)] -def _k128_covers(tokens: int) -> bool: - """gemm2K128's gate at tile N 8: one CTA per SM, at most 15 units (token - tile x 128 output features) each, at most 64 token tiles.""" +def _small_k_covers(tokens: int, inter: int = 128) -> bool: + """Owned FC2 gate: K128 allows 15 units/64 tiles, K256 allows 29 units/128 tiles.""" tiles = min(tokens * TOP_K, EXPERTS) sms = torch.cuda.get_device_properties(0).multi_processor_count units = tiles * HIDDEN // 128 - return tiles <= 64 and -(-units // sms) <= 15 + return tiles <= (64 if inter == 128 else 128) and -(-units // sms) <= ( + 15 if inter == 128 else 29 + ) # ── tests ───────────────────────────────────────────────────────────────── @@ -401,7 +406,7 @@ def _mutation_case_args(fused_router, do_finalize, caller_output): hidden_states_scale=case["xsf"], fused_router=fused_router, workspace=make_fused_router_workspace(DEV) if fused_router else None, - gemm2_k128=True, + specialized_gemm2=True, do_finalize=do_finalize, routing_logits=case["logits"], output=torch.zeros(2, HIDDEN, device=DEV, dtype=torch.bfloat16) if caller_output else None, @@ -506,9 +511,18 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 certifying(op) spec = cell.spec geometry = (spec["hidden"], spec["intermediate"], spec["num_experts"], spec["top_k"]) - assert geometry == (HIDDEN, INTER, EXPERTS, TOP_K), geometry + assert geometry in ( + (HIDDEN, INTER, EXPERTS, TOP_K), + (HIDDEN, 256, EXPERTS, TOP_K), + (HIDDEN, 512, EXPERTS, TOP_K), + ), geometry assert spec["routed_scaling_factor"] == ROUTED_SCALE and spec["local_expert_offset"] == 0 - case = _case(spec["tokens"], seed=zlib.crc32(cell.why.encode()) % 2**31, clear_margin=True) + case = _case( + spec["tokens"], + seed=zlib.crc32(cell.why.encode()) % 2**31, + clear_margin=True, + inter=spec["intermediate"], + ) workspace = make_fused_router_workspace(DEV) if spec["fused_router"] else None if spec["late_join"]: side = torch.cuda.Stream() @@ -522,11 +536,11 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 fused_router=spec["fused_router"], workspace=workspace, event=ready, - gemm2_k128=spec["gemm2_k128"], + specialized_gemm2=spec["specialized_gemm2"], ) torch.cuda.synchronize() - stack = _stack() + stack = _stack(spec["intermediate"]) ref, _, ref_weights = op.reference( routing_logits=None if spec["fused_router"] else case["logits"], routing_bias=case["routing_bias"], @@ -568,27 +582,41 @@ def test_fused_router_matches_the_router_gemm(tokens) -> None: case.update(x=x, router_weight=weight) logits = torch.ops.trtllm.dsv3_router_gemm_op(x, weight.t(), None, torch.float32) workspace = make_fused_router_workspace(DEV) - common = dict(hidden_states=case["xq"], hidden_states_scale=case["xsf"], gemm2_k128=False) + common = dict( + hidden_states=case["xq"], hidden_states_scale=case["xsf"], specialized_gemm2=False + ) fused = _call(case, fused_router=True, workspace=workspace, **common) routed = _call(case, fused_router=False, routing_logits=logits, **common) _assert_same(fused, routed, f"{tokens} rows") -@pytest.mark.parametrize("tokens", (1, 8)) -def test_gemm2_k128_matches_the_cubin(tokens) -> None: - """FC2 at K = 128 writes the Gemm2 cubin's rows, bitwise, and runs where - its gate says it does.""" - case = _case(tokens, seed=4200 + tokens) - common = dict(hidden_states=case["xq"], hidden_states_scale=case["xsf"], fused_router=False) - k128 = _call(case, gemm2_k128=True, **common) - cubin = _call(case, gemm2_k128=False, **common) - _assert_same(k128, cubin, f"{tokens} rows") - - names = _kernels(lambda: _call(case, gemm2_k128=True, **common)) - ran = any("gemm2K128Kernel" in n for n in names) - assert ran == _k128_covers(tokens), (names, _k128_covers(tokens)) - names = _kernels(lambda: _call(case, gemm2_k128=False, **common)) - assert not any("gemm2K128Kernel" in n for n in names), names +@pytest.mark.parametrize("tokens,inter", ((1, 128), (8, 128), (4, 256), (8, 256))) +@pytest.mark.parametrize("do_finalize", (False, True)) +def test_specialized_gemm2_matches_the_cubin(tokens, inter, do_finalize) -> None: + """Owned FC2 matches generic FC2, including in-place caller-output finalization.""" + case = _case(tokens, seed=4200 + tokens, inter=inter) + common = dict( + hidden_states=case["xq"], + hidden_states_scale=case["xsf"], + fused_router=False, + do_finalize=do_finalize, + ) + output = torch.empty(tokens, HIDDEN, dtype=torch.bfloat16, device=DEV) if do_finalize else None + specialized = _call(case, specialized_gemm2=True, output=output, **common) + cubin = _call(case, specialized_gemm2=False, **common) + if output is None: + _assert_same(specialized, cubin, f"{tokens} rows, K{inter}") + else: + # The mutated caller output cannot alias a custom-op return. + assert len(specialized) == 1 and specialized[0].numel() == 0 + _assert_same([output], cubin, f"{tokens} rows, K{inter}, caller output") + allocated = _call(case, specialized_gemm2=True, **common) + _assert_same(allocated, cubin, f"{tokens} rows, K{inter}, allocated output") + marker = f"gemm2K{inter}Kernel" + names = _kernels(lambda: _call(case, specialized_gemm2=True, **common)) + assert any(marker in name for name in names) == _small_k_covers(tokens, inter), names + names = _kernels(lambda: _call(case, specialized_gemm2=False, **common)) + assert not any("gemm2K128Kernel" in name or "gemm2K256Kernel" in name for name in names), names def test_c1_call_launches_three_kernels() -> None: @@ -604,13 +632,13 @@ def test_c1_call_launches_three_kernels() -> None: hidden_states_scale=case["xsf"], fused_router=True, workspace=workspace, - gemm2_k128=True, + specialized_gemm2=True, routing_logits=shape_only, ) ) assert len(names) == 3, names assert "routingFusedRouterKernel" in names[0], names - assert ("gemm2K128Kernel" in names[2]) == _k128_covers(8), names + assert ("gemm2K128Kernel" in names[2]) == _small_k_covers(8), names def test_late_join_waits_for_the_input() -> None: @@ -626,7 +654,7 @@ def test_late_join_waits_for_the_input() -> None: hidden_states_scale=case["xsf"], fused_router=True, workspace=workspace, - gemm2_k128=True, + specialized_gemm2=True, ) buffers = (stale["xq"].clone(), stale["xsf"].clone()) side = torch.cuda.Stream() @@ -638,18 +666,19 @@ def test_late_join_waits_for_the_input() -> None: fused_router=True, workspace=workspace, event=ready, - gemm2_k128=True, + specialized_gemm2=True, ) torch.cuda.synchronize() _assert_same(late, expected, "late join") -def test_c1_call_replays_in_a_cuda_graph() -> None: +@pytest.mark.parametrize("tokens,inter", [(8, 128), (12, 128), (18, 128), (4, 256), (8, 256)]) +def test_fused_call_replays_in_a_cuda_graph(tokens, inter) -> None: """The c1 step captured as the target runs it -- input quantized on a side stream, joined inside the op -- replays the eager result on new inputs, and every launch leaves the router workspace ready for the next one.""" - cases = [_case(8, seed=4500 + i) for i in range(3)] - static = {k: v.clone() for k, v in cases[0].items()} + cases = [_case(tokens, seed=4500 + i, inter=inter) for i in range(3)] + static = {k: v.clone() if isinstance(v, torch.Tensor) else v for k, v in cases[0].items()} workspace = make_fused_router_workspace(DEV) side = torch.cuda.Stream() @@ -662,7 +691,7 @@ def step(): fused_router=True, workspace=workspace, event=ready, - gemm2_k128=True, + specialized_gemm2=tokens <= 8, routing_logits=static["logits"], ) @@ -675,7 +704,8 @@ def step(): launches = 1 for case in cases: for key, value in case.items(): - static[key].copy_(value) + if isinstance(value, torch.Tensor): + static[key].copy_(value) graph.replay() launches += 1 torch.cuda.synchronize() @@ -685,7 +715,7 @@ def step(): hidden_states_scale=case["xsf"], fused_router=True, workspace=workspace, - gemm2_k128=True, + specialized_gemm2=tokens <= 8, ) launches += 1 torch.cuda.synchronize() @@ -697,10 +727,10 @@ def step(): assert bool((current == -(2**31)).all()), "the next launch's exchange buffer is not idle" -def test_fused_router_refuses_more_than_8_rows() -> None: - """Past 8 rows the fused router is a loud error; the caller routes from +def test_fused_router_refuses_more_than_24_rows() -> None: + """Past 24 rows the fused router is a loud error; the caller routes from the router GEMM's logits instead (the 32-row cell).""" - case = _case(9, seed=4600) + case = _case(25, seed=4600) with pytest.raises(RuntimeError, match="fused router"): _call( case, @@ -708,7 +738,7 @@ def test_fused_router_refuses_more_than_8_rows() -> None: hidden_states_scale=case["xsf"], fused_router=True, workspace=make_fused_router_workspace(DEV), - gemm2_k128=True, + specialized_gemm2=True, ) @@ -726,7 +756,7 @@ def test_guard_refuses_a_strided_router_input() -> None: hidden_states_scale=case["xsf"], fused_router=True, workspace=make_fused_router_workspace(DEV), - gemm2_k128=True, + specialized_gemm2=True, ) @@ -743,7 +773,7 @@ def test_guard_refuses_a_zero_filled_workspace() -> None: hidden_states_scale=case["xsf"], fused_router=True, workspace=workspace, - gemm2_k128=True, + specialized_gemm2=True, ) @@ -759,5 +789,38 @@ def test_guard_refuses_an_unrecorded_event() -> None: fused_router=True, workspace=make_fused_router_workspace(DEV), event=torch.cuda.Event(), - gemm2_k128=True, + specialized_gemm2=True, ) + + +@pytest.mark.parametrize("tokens", [9, 12, 16, 17, 18, 24]) +def test_wide_fused_router_matches_separate_routing(tokens) -> None: + """Exact-grid router operands isolate routing/tile semantics from GEMM rounding.""" + case = _case(tokens, seed=6100 + tokens, clear_margin=True) + common = dict( + hidden_states=case["xq"], hidden_states_scale=case["xsf"], specialized_gemm2=False + ) + fused = _call(case, fused_router=True, workspace=make_fused_router_workspace(DEV), **common) + separate = _call(case, fused_router=False, **common) + torch.cuda.synchronize() + # Physical tile order may differ: compare each token's finalized value and routing weights. + _compare(_finalize(fused), _finalize(separate)) + assert torch.equal(fused[1].sort(dim=-1).values, separate[1].sort(dim=-1).values) + + +def test_fused_router_workspace_alternates_small_and_wide_calls() -> None: + """Both kernels rotate/reset the same16388-word workspace across changing row widths.""" + workspace = make_fused_router_workspace(DEV) + for launch, tokens in enumerate((8, 12, 1, 18, 6, 24, 8), start=1): + case = _case(tokens, seed=6200 + tokens, clear_margin=True) + common = dict( + hidden_states=case["xq"], hidden_states_scale=case["xsf"], specialized_gemm2=False + ) + fused = _call(case, fused_router=True, workspace=workspace, **common) + separate = _call(case, fused_router=False, **common) + torch.cuda.synchronize() + _compare(_finalize(fused), _finalize(separate)) + epoch = int(workspace[_EPOCH_WORD].item()) + assert epoch == launch + current = workspace[(epoch & 1) * _BUFFER_WORDS : ((epoch & 1) + 1) * _BUFFER_WORDS] + assert bool((current == -(2**31)).all()) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/quantization/test_modeling_v2_glm_fp4_quantize.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/quantization/test_modeling_v2_glm_fp4_quantize.py similarity index 96% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/quantization/test_modeling_v2_glm_fp4_quantize.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/quantization/test_modeling_v2_glm_fp4_quantize.py index 305e9771c770..f320ae21708f 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/quantization/test_modeling_v2_glm_fp4_quantize.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/quantization/test_modeling_v2_glm_fp4_quantize.py @@ -12,7 +12,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( fp4_quantize as shared_op, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.quantization.fp4_quantize import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.quantization.fp4_quantize import ( # noqa: E501 fp4_quantize as op, ) from tensorrt_llm._utils import get_sm_version diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/_draft_step_op_matrix.py similarity index 65% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/_draft_step_op_matrix.py index 096da8663092..93503df0b8d7 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/_draft_step_op_matrix.py @@ -22,6 +22,7 @@ tensorrt_llm, and the ranks reach the target's ops by absolute import. """ +import contextlib import os import signal import socket @@ -30,7 +31,8 @@ import time import zlib from pathlib import Path -from typing import Any, Dict, List +from types import SimpleNamespace +from typing import Any, Callable, Dict, List import torch import torch.distributed as dist @@ -107,6 +109,9 @@ def _state(spec: dict, seed: int) -> Dict[str, Any]: vocab = spec["vocab_shard"] * WORLD num_accepted = torch.randint(1, draft_len + 2, (rows,), generator=g, device="cuda").int() batch_indices = torch.randperm(rows, generator=g, device="cuda").int() + if WORLD == 8 and draft_len == 3 and rows == 2: + num_accepted.copy_(torch.tensor([1, 4], dtype=torch.int32, device="cuda")) + batch_indices.copy_(torch.tensor([1, 0], dtype=torch.int32, device="cuda")) state: Dict[str, Any] = dict( draft_tokens=torch.randint(0, vocab, (rows, draft_len), generator=g, device="cuda").int(), draft_step=step, @@ -326,7 +331,9 @@ def check_glue_cells() -> int: for index, cell in enumerate(glue_op.CELLS): spec = cell.spec seed = zlib.crc32(cell.why.encode()) % 2**31 - if spec["glue"] and embed_table is None: + if spec["glue"] and ( + embed_table is None or embed_table.shape[0] != spec["vocab_shard"] * WORLD + ): embed_table = torch.randn( spec["vocab_shard"] * WORLD, HIDDEN, @@ -344,24 +351,44 @@ def check_glue_cells() -> int: return len(glue_op.CELLS) -def _draft_loop_specs(draft_len: int = 7): - """The target's c1 draft loop: table steps 0..K-2, then the glue for step K-1.""" +def _draft_loop_specs(draft_len: int | None = None, rows: int = 1): + """TP8 K3/K7 or TP16 K7: table steps 0..K-2, then final sample/glue.""" + if draft_len is None: + draft_len = 3 if WORLD == 8 else 7 + shard = 19360 if WORLD == 8 else 9680 table_steps = [] for step in range(draft_len - 1): table_steps.append( - table_op.CELLS[0].spec if step == 0 else dict(table_op.CELLS[1].spec, step=step) + dict( + table_op.CELLS[0 if step == 0 else 1].spec, + rows=rows, + draft_len=draft_len, + vocab_shard=shard, + step=step, + ) ) - glue_step = dict(glue_op.CELLS[0].spec, step=draft_len - 1) + glue_step = dict( + glue_op.CELLS[0].spec, rows=rows, draft_len=draft_len, vocab_shard=shard, step=draft_len - 1 + ) return table_steps, glue_step def check_draft_loop_eager_and_graph() -> None: + shapes = ((2, 3), (1, 3), (1, 7)) if WORLD == 8 else ((1, 7),) + replays = [_check_draft_loop_eager_and_graph(rows, k) for rows, k in shapes] + if WORLD == 8: + for repeat in range(3): + for replay in replays: + replay(repeat) + + +def _check_draft_loop_eager_and_graph(rows: int, draft_len: int) -> Callable[[int], None]: """Seven consecutive calls per decode step over two workspaces, eagerly and as graph replays. The Lamport flags rotate on every call and each call empties the buffer of the previous one, so the exchange must stay correct across steps, decode steps and replays of a captured loop. """ - table_steps, glue_step = _draft_loop_specs() + table_steps, glue_step = _draft_loop_specs(draft_len=draft_len, rows=rows) calls = [] for step, spec in enumerate(table_steps): seed = 1000 + step @@ -405,7 +432,8 @@ def reset_and_expect(): with torch.cuda.graph(graph, stream=side): run_loop() torch.cuda.current_stream().wait_stream(side) - for repeat in range(3): + + def replay(repeat): expected = reset_and_expect() graph.replay() torch.cuda.synchronize() @@ -413,6 +441,238 @@ def reset_and_expect(): names = (_TABLE_OUTPUTS if kind == "table" else _GLUE_OUTPUTS) + _STATE_OUTPUTS _compare(exp, _written(kwargs, names), f"graph replay {repeat} {kind}") + for repeat in range(3): + replay(repeat) + return replay + + +def check_tp8_worker_feedback_eager_and_graph() -> None: + """Real K3/K7 worker/table feedback, with deterministic attention/head substitutes. + + Native gather, table and final sample/glue run on the production workspaces. + The synthetic attention consumes both projected tensors; the next hidden half + consumes that output, so independent per-step fixtures cannot satisfy this check. + """ + if WORLD != 8: + return + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency._mtp_worker import ( # noqa: E501 + MTPDraftWorker, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency._small_batch_draft import ( # noqa: E501 + SmallBatchDraftKernels, + ) + + shard = 19360 + tables = _LazyTables(812, shard) + # Materialize local and reference winner tables before capture. + local_table, winner_table = tables[RANK], tables[WORLD - 1] + graphs = [] + + for batch, draft_len, accepted in ((2, 3, [1, 4]), (1, 3, [1]), (1, 7, [1]), (1, 7, [8])): + window = draft_len + 1 + rows = batch * window + order = torch.tensor(list(reversed(range(batch))), dtype=torch.int32, device="cuda") + counts = torch.tensor(accepted, dtype=torch.int32, device="cuda") + ids = torch.arange(rows, dtype=torch.int32, device="cuda") + positions = torch.arange(rows, dtype=torch.int64, device="cuda") + 1000 + hidden = torch.randn(rows, HIDDEN, generator=_generator(batch), device="cuda").to( + torch.bfloat16 + ) + seq = torch.full((batch,), window, dtype=torch.int32, device="cuda") + kv_initial = torch.arange(batch, dtype=torch.int32, device="cuda") * 37 + 1000 + kv = kv_initial.clone() + meta = SimpleNamespace( + _seq_lens=torch.full((batch,), window, dtype=torch.int32), + seq_lens_cuda=seq, + kv_lens_cuda=kv, + num_contexts=0, + use_spec_decoding=True, + kv_cache_manager="target", + updates=[], + ) + meta.set_in_mtp_draft_loop = lambda value, m=meta: setattr(m, "in_mtp_draft_loop", value) + meta.set_skip_topk = lambda value, m=meta: setattr(m, "indexer_skip_topk", value) + meta.set_mtp_num_accepted = lambda value, m=meta: setattr(m, "mtp_num_accepted", value) + meta.on_update = lambda: None + meta.update_for_spec_dec = lambda m=meta, s=seq, k=kv: m.updates.append( + (s.clone(), k.clone()) + ) + helper = object.__new__(SmallBatchDraftKernels) + helper.use_table, helper.hidden_size = True, HIDDEN + helper.DRAFT_LEN = draft_len + helper.decode = SimpleNamespace(mapping=SimpleNamespace(tp_size=WORLD)) + helper.sample_exchange, helper.table_exchange = WORKSPACES["glue"], WORKSPACES["table"] + helper.token_table, helper.half_done = local_table, torch.cuda.Event() + helper.layer = SimpleNamespace( + input_layernorm=SimpleNamespace( + weight=torch.ones(HIDDEN, dtype=torch.bfloat16, device="cuda"), variance_epsilon=EPS + ) + ) + helper.supports = ( + lambda metadata, first, selected, window=window: first == selected * window + ) + values = [] + seen = [] + + def head(value, *, next_step=False, h=helper, values=values, b=batch): + step = len(values) + values.append(value) + logits = value.new_full((b, shard), -100) + token = (torch.arange(b, device="cuda") * 101 + step * 37 + RANK * 7).long() + logits.scatter_(1, token[:, None], RANK + 1) + half = ( + value[:, RANK * (HIDDEN // WORLD) : (RANK + 1) * (HIDDEN // WORLD)] + .float() + .contiguous() + ) + h.half_done.record() + return logits, half if next_step else None + + helper.head = head + + @contextlib.contextmanager + def draft_context(metadata, manager): + original = metadata.kv_cache_manager + metadata.kv_cache_manager = manager + try: + yield metadata + finally: + metadata.kv_cache_manager = original + + def forward(model, inputs, spec, step, seen=seen): + seen.append((inputs["hidden_states"].shape[0], inputs["position_ids"].clone())) + if step == 0: + value, residual = inputs["row_select"]( + inputs["hidden_states"], inputs["hidden_states"] + 10 + ) + else: + value, residual = inputs["eh_proj_output"] + assert "eh_proj_input" not in inputs and "row_select" not in inputs + return (value.float() + 0.01 * residual.float()).to(torch.bfloat16), None + + worker = SimpleNamespace( + _drafted_rows=None, draft_kv_cache_context=draft_context, _run_draft_forward=forward + ) + spec = SimpleNamespace(runtime_draft_len=draft_len, batch_indices_cuda=order) + inputs = dict( + input_ids=ids, + position_ids=positions, + hidden_states=hidden, + attn_metadata=meta, + spec_metadata=spec, + ) + + def run( + w=worker, + h=helper, + ins=inputs, + m=meta, + sp=spec, + b=batch, + accepted=counts, + v=values, + s=seen, + ): + v.clear() + s.clear() + m.updates.clear() + drafted = MTPDraftWorker._forward_small_batch_draft_loop( + w, h, ins, m, sp, object(), "draft", b, accepted + ) + return drafted, w._drafted_rows[1], list(v), list(s), list(m.updates) + + def reset(m=meta, seq=seq, kv=kv, initial=kv_initial, window=window): + m._seq_lens.fill_(window) + seq.fill_(window) + kv.copy_(initial) + m.use_spec_decoding = True + + def check( + result, + b=batch, + order=order, + counts=counts, + hidden=hidden, + positions=positions, + initial=kv_initial, + m=meta, + draft_len=draft_len, + window=window, + ): + drafted, next_tokens, step_values, seen_rows, updates = result + selected = order.long() * window + counts.long() - 1 + value = (hidden[selected].float() + 0.01 * (hidden[selected] + 10).float()).to( + torch.bfloat16 + ) + expected_tokens = torch.stack( + [ + torch.arange(b, device="cuda", dtype=torch.int32) * 101 + + step * 37 + + (WORLD - 1) * 7 + + (WORLD - 1) * shard + for step in range(draft_len) + ], + dim=1, + ) + _assert_exact("draft_tokens", drafted, expected_tokens, "connected worker") + _assert_exact("accepted", next_tokens[:, 0], selected.int(), "connected worker") + _assert_exact("next drafts", next_tokens[:, 1:], drafted, "connected worker") + assert [item[0] for item in seen_rows] == [b * window] + [b] * (draft_len - 1) + for step in range(draft_len): + assert_within_ulp(step_values[step], value, element_ulp=2.0, rms_ulp=0.5) + if step: + _assert_exact( + "positions", + seen_rows[step][1], + positions[selected] + step, + "connected worker", + ) + _assert_exact( + "seq_lens", updates[step][0], torch.ones_like(counts), "connected worker" + ) + _assert_exact( + "kv_lens", + updates[step][1], + initial - (draft_len - counts) + step, + "connected worker", + ) + if step < draft_len - 1: + residual = ( + winner_table[(expected_tokens[:, step] % shard).long()].float() + + value.float() + ).to(torch.bfloat16) + normed = ( + residual.float() + * torch.rsqrt(residual.float().square().mean(-1, keepdim=True) + EPS) + ).to(torch.bfloat16) + value = (normed.float() + 0.01 * residual.float()).to(torch.bfloat16) + assert m.kv_cache_manager == "target" and not m.in_mtp_draft_loop + assert not m.indexer_skip_topk and m.mtp_num_accepted is None + + reset() + eager = run() + torch.cuda.synchronize() + check(eager) + reset() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + result = run() + # Retain captured input buffers and the helper's CUDA events for every graph. + graphs.append((graph, reset, check, result, hidden, positions, helper, inputs)) + + for repeat in range(3): + for graph, reset, check, result, hidden, positions, _helper, _inputs in graphs: + hidden.copy_( + torch.randn( + hidden.shape, generator=_generator(900 + repeat, hidden.shape[0]), device="cuda" + ).to(torch.bfloat16) + ) + positions.add_(17) + reset() + graph.replay() + torch.cuda.synchronize() + check(result) + def check_guards_refuse() -> None: """The guards refuse their silent-wrong inputs on every rank before anything is launched.""" @@ -449,39 +709,22 @@ def check_guards_refuse() -> None: raise AssertionError("the glue guard accepted a wider next_new_tokens") -def _rank_env(): - env = os.environ - rank = int(env.get("RANK", env.get("SLURM_PROCID", "0"))) - world = int(env.get("WORLD_SIZE", env.get("SLURM_NTASKS", "1"))) - local_rank = int(env.get("LOCAL_RANK", env.get("SLURM_LOCALID", str(rank)))) - return rank, world, local_rank - - def _run_one_rank(entry: str, expected_cells: int) -> int: """Body of one rank: build the workspaces, run the entry's checks, exit hard on a failure.""" global RANK, WORLD, HOST_GROUP, table_op, glue_op, validating, assert_within_ulp - os.environ.setdefault("TLLM_DISABLE_MPI", "1") - rank, world, local_rank = _rank_env() - os.environ.setdefault("RANK", str(rank)) - os.environ.setdefault("WORLD_SIZE", str(world)) - torch.cuda.set_device(local_rank) - dist.init_process_group( - backend="cuda:nccl,cpu:gloo", - init_method="env://", - rank=rank, - world_size=world, - device_id=torch.device("cuda", local_rank), - ) - RANK, WORLD = rank, world + sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "comm")) + from _mnnvl_ranks import init_rank + + RANK, WORLD, mapping = init_rank() HOST_GROUP = dist.new_group(backend="gloo") - from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.spec import ( # noqa: E501 _draft_exchange, ) - from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec.draft_ehproj_table import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.spec.draft_ehproj_table import ( # noqa: E501 glm_draft_ehproj_table, ) - from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec.draft_sample_glue import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.spec.draft_sample_glue import ( # noqa: E501 glm_draft_sample_glue, ) @@ -494,34 +737,27 @@ def _run_one_rank(entry: str, expected_cells: int) -> int: validating, assert_within_ulp = _validating_cm, _within - hosts: List[Any] = [None] * WORLD - dist.all_gather_object(hosts, socket.gethostname(), group=HOST_GROUP) - fabric = len(set(hosts)) > 1 or os.environ.get("TRTLLM_FORCE_MNNVL_AR", "0") == "1" + from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce + + if WORLD == 8: + assert MNNVLAllReduce.is_mnnvl(mapping, torch.bfloat16), ( + "TP8 requires the actual MNNVL capability" + ) assert _draft_exchange.ehproj_table_supports_tp(WORLD), ( f"the table op does not serve {WORLD} ranks" ) - max_rows = _draft_exchange.MAX_ROWS - WORKSPACES["table"] = _draft_exchange.DraftExchangeWorkspace( - max_rows * WORLD * _draft_exchange.ehproj_table_slot_bytes(WORLD), - WORLD, - RANK, - local_rank, - fabric, - dist.group.WORLD, - ) - WORKSPACES["glue"] = _draft_exchange.DraftExchangeWorkspace( - max_rows * WORLD * _draft_exchange.sample_glue_slot_bytes(WORLD), - WORLD, - RANK, - local_rank, - fabric, - dist.group.WORLD, - ) + WORKSPACES["table"] = _draft_exchange.DraftExchangeWorkspace.for_ehproj_table(mapping) + WORKSPACES["glue"] = _draft_exchange.DraftExchangeWorkspace.for_sample_glue(mapping) + assert WORKSPACES["table"].buffer.data_ptr() != WORKSPACES["glue"].buffer.data_ptr() checks = { "draft_ehproj_table": (check_table_cells,), "draft_sample_glue": (check_glue_cells,), - }[entry] + (check_draft_loop_eager_and_graph, check_guards_refuse) + }[entry] + ( + check_draft_loop_eager_and_graph, + check_tp8_worker_feedback_eager_and_graph, + check_guards_refuse, + ) for check in checks: try: ran = check() diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_rank_job.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/_draft_step_rank_job.py similarity index 100% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/_draft_step_rank_job.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/_draft_step_rank_job.py diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_ehproj_table.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_draft_ehproj_table.py similarity index 80% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_ehproj_table.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_draft_ehproj_table.py index b536ec7d560e..de235aab5f8b 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_ehproj_table.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_draft_ehproj_table.py @@ -9,14 +9,16 @@ import _draft_step_rank_job import pytest +import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec.draft_ehproj_table import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.spec.draft_ehproj_table import ( # noqa: E501 glm_draft_ehproj_table as op, ) from tensorrt_llm._utils import get_sm_version pytestmark = pytest.mark.skipif( - get_sm_version() != 107, reason="the GLM-5.2 tp16 target's ops are certified on sm_107" + get_sm_version() != 107 or torch.cuda.device_count() < _draft_step_rank_job.WORLD_SIZE, + reason="requires SM107 and at least four visible GPUs", ) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_exchange.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_draft_exchange.py similarity index 98% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_exchange.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_draft_exchange.py index 07348dd310e8..66a44c65236f 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_exchange.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_draft_exchange.py @@ -9,7 +9,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec import ( +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.spec import ( # noqa: E501 _draft_exchange as exchange, ) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py similarity index 98% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py index 3e211989b671..7742def831b3 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_draft_glue_rmsnorm_cat.py @@ -7,7 +7,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec.draft_glue_rmsnorm_cat import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.spec.draft_glue_rmsnorm_cat import ( # noqa: E501 glm_draft_glue_rmsnorm_cat as op, ) from tensorrt_llm._utils import get_sm_version diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_sample_glue.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_draft_sample_glue.py similarity index 80% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_sample_glue.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_draft_sample_glue.py index faf00183c753..956700cddd46 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_sample_glue.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_draft_sample_glue.py @@ -9,14 +9,16 @@ import _draft_step_rank_job import pytest +import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec.draft_sample_glue import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.spec.draft_sample_glue import ( # noqa: E501 glm_draft_sample_glue as op, ) from tensorrt_llm._utils import get_sm_version pytestmark = pytest.mark.skipif( - get_sm_version() != 107, reason="the GLM-5.2 tp16 target's ops are certified on sm_107" + get_sm_version() != 107 or torch.cuda.device_count() < _draft_step_rank_job.WORLD_SIZE, + reason="requires SM107 and at least four visible GPUs", ) diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_token_table.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_draft_token_table.py similarity index 95% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_token_table.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_draft_token_table.py index 54b0750306bf..98ae69c08739 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_token_table.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_draft_token_table.py @@ -6,7 +6,7 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import true_fp32_matmul -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec._draft_token_table import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.spec._draft_token_table import ( # noqa: E501 build_draft_token_table, ) from tensorrt_llm._utils import get_sm_version @@ -22,7 +22,7 @@ EPS = 1e-5 -@pytest.mark.parametrize("tp_size,tp_rank", [(16, 0), (16, 11), (4, 3)]) +@pytest.mark.parametrize("tp_size,tp_rank", [(16, 0), (16, 11), (4, 3), (8, 0), (8, 7)]) def test_table_is_the_token_half_of_eh_proj(tp_size, tp_rank) -> None: """Each table row is eh_proj's embedding half of one token of the rank's shard. diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_gather_last_accepted_rows.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_gather_last_accepted_rows.py similarity index 97% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_gather_last_accepted_rows.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_gather_last_accepted_rows.py index 3edc7f4423e5..e0243e86e299 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_gather_last_accepted_rows.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_gather_last_accepted_rows.py @@ -7,7 +7,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec.gather_last_accepted_rows import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.spec.gather_last_accepted_rows import ( # noqa: E501 glm_gather_last_accepted_rows as op, ) from tensorrt_llm._utils import get_sm_version diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_small_batch_mtp.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_small_batch_mtp.py new file mode 100644 index 000000000000..9100a1201c63 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_small_batch_mtp.py @@ -0,0 +1,605 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU checks of MTP row ordering, metadata transitions and failure cleanup. + +The native sample/glue is replaced by its existing mathematical reference; +these tests exercise the real bridge around it, not GPU kernel numerics. +""" + +import contextlib +import inspect +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency import ( + _mtp_worker as worker_module, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency import ( + _small_batch_draft as small, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency._mtp_worker import ( # noqa: E501 + MTPDraftWorker, +) + +pytestmark = pytest.mark.cpu_only + + +@pytest.fixture +def reference_ops(monkeypatch): + gather = small.glm_gather_last_accepted_rows + sample = small.glm_draft_sample_glue + signature = inspect.signature(sample.reference) + table = small.glm_draft_ehproj_table + table_signature = inspect.signature(table.reference) + calls = [] + + def sample_reference(logits, *args): + # Rank zero wins; other shards still participate in the global argmax. + bound = signature.bind(logits, *args) + bound.arguments["logits"] = [logits] + [ + torch.full_like(logits, -float("inf")) for _ in range(bound.arguments["group_size"] - 1) + ] + outputs = sample.reference(*bound.args, **bound.kwargs) + for name, value in outputs.items(): + bound.arguments[name].copy_(value) + calls.append(("sample", bound.arguments["draft_step"], outputs)) + + def table_reference(logits, *args): + bound = table_signature.bind(logits, *args) + world = bound.arguments["group_size"] + bound.arguments["logits"] = [logits] + [ + torch.full_like(logits, -float("inf")) for _ in range(world - 1) + ] + bound.arguments["hidden_half"] = [bound.arguments["hidden_half"]] * world + bound.arguments["token_table"] = [bound.arguments["token_table"]] * world + outputs = table.reference(*bound.args, **bound.kwargs) + for name, value in outputs.items(): + bound.arguments[name].copy_(value) + calls.append(("table", bound.arguments["draft_step"], outputs)) + + monkeypatch.setattr(small, "glm_gather_last_accepted_rows", gather.reference) + monkeypatch.setattr(small, "glm_draft_sample_glue", sample_reference) + monkeypatch.setattr(small, "glm_draft_ehproj_table", table_reference) + monkeypatch.setattr( + torch.cuda, "current_stream", lambda: SimpleNamespace(wait_event=lambda event: None) + ) + return calls + + +class _Metadata: + def __init__(self, batch, k=5): + self._seq_lens = torch.full((batch,), k + 1, dtype=torch.int32) + self.seq_lens_cuda = self._seq_lens.clone() + self.kv_lens_cuda = torch.arange(batch, dtype=torch.int32) + 100 + self.num_contexts = 0 + self.num_tokens = batch * (k + 1) + self.kv_cache_manager = "target" + self.use_spec_decoding = True + self.in_mtp_draft_loop = False + self.indexer_skip_topk = False + self.mtp_num_accepted = None + self.updates = [] + + def set_in_mtp_draft_loop(self, value): + self.in_mtp_draft_loop = value + + def set_skip_topk(self, value): + self.indexer_skip_topk = value + + def set_mtp_num_accepted(self, value): + self.mtp_num_accepted = value + + def on_update(self): + self.num_tokens = int(self._seq_lens.sum()) + + def update_for_spec_dec(self): + self.updates.append((self.seq_lens_cuda.clone(), self.kv_lens_cuda.clone())) + + +def _loop_fixture(batch, fail_step=None, eligible=True, tp_size=4, k=None, accepted=None): + k = (3 if tp_size == 8 else 5) if k is None else k + hidden = 8 if tp_size == 8 else 4 + rows = batch * (k + 1) + metadata = _Metadata(batch, k) + # Permutation and heterogeneous accepted counts expose double selection and + # selection of the last verify row instead of the last accepted row. + order = torch.tensor( + [2, 0, 1][:batch] if batch == 3 else list(reversed(range(batch))), dtype=torch.int32 + ) + counts = torch.tensor( + [1, k + 1, 3][:batch] if accepted is None else accepted, dtype=torch.int32 + ) + spec = SimpleNamespace(runtime_draft_len=k, batch_indices_cuda=order) + ids = torch.arange(rows, dtype=torch.int32) + positions = torch.arange(rows, dtype=torch.int64) + 1000 + states = torch.arange(rows, dtype=torch.bfloat16).unsqueeze(1).expand(rows, hidden).clone() + inputs = dict( + input_ids=ids, + position_ids=positions, + hidden_states=states, + attn_metadata=metadata, + spec_metadata=spec, + ) + helper = object.__new__(small.SmallBatchDraftKernels) + helper.use_table = tp_size == 8 + helper.DRAFT_LEN = k + helper.hidden_size = hidden + helper.window = (0, 2) + helper.decode = SimpleNamespace(mapping=SimpleNamespace(tp_size=tp_size)) + helper.sample_exchange = SimpleNamespace(buffer=torch.empty(0), flags=torch.empty(0)) + helper.embed_table = torch.ones(32, hidden, dtype=torch.bfloat16) + helper.enorm_weight = torch.ones(hidden, dtype=torch.bfloat16) + helper.hnorm_weight = torch.ones(hidden, dtype=torch.bfloat16) + helper.layer = SimpleNamespace( + enorm=SimpleNamespace(variance_epsilon=1e-5), + input_layernorm=SimpleNamespace(weight=torch.ones(hidden), variance_epsilon=1e-5), + ) + helper.table_exchange = SimpleNamespace(buffer=torch.empty(0), flags=torch.empty(0)) + helper.token_table = torch.arange(8 * hidden, dtype=torch.bfloat16).reshape(8, hidden) + 1 + helper.half_done = object() + helper.head_inputs = [] + helper.projected_inputs = [] + calls = [] + preflights = [] + + def supports(meta, first_rows, selected_rows): + preflights.append((meta.kv_cache_manager, first_rows, selected_rows)) + return eligible + + def head(value, *, next_step=False): + helper.head_inputs.append(value.clone()) + logits = torch.zeros(batch, 8, dtype=torch.bfloat16) + logits.scatter_(1, value[:, :1].long() % 8, 1) + return logits, value[:, :1].float() if helper.use_table and next_step else None + + helper.supports = supports + helper.head = head + + @contextlib.contextmanager + def draft_context(meta, manager): + saved = meta.kv_cache_manager + meta.kv_cache_manager = manager + try: + yield meta + finally: + meta.kv_cache_manager = saved + + def run_forward(draft_model, step_inputs, spec_metadata, step): + calls.append( + ( + step, + step_inputs["hidden_states"].shape[0], + step_inputs["position_ids"].clone(), + metadata.kv_lens_cuda.clone(), + metadata.indexer_skip_topk, + ) + ) + if step == fail_step: + raise RuntimeError("injected draft failure") + projected = step_inputs.get("eh_proj_output") + helper.projected_inputs.append(projected) + value = step_inputs["hidden_states"] if projected is None else projected[0] + select = step_inputs.get("row_select") + if select is not None: + value, _ = select(value, value + 100) + return value, None + + worker = SimpleNamespace( + _drafted_rows=None, draft_kv_cache_context=draft_context, _run_draft_forward=run_forward + ) + return worker, helper, inputs, metadata, spec, counts, calls, preflights + + +@pytest.mark.parametrize("batch", [1, 2, 3]) +def test_loop_selects_once_and_advances_metadata_after_every_step(reference_ops, batch): + worker, helper, inputs, meta, spec, counts, calls, preflights = _loop_fixture(batch) + original_positions = inputs["position_ids"].clone() + original_kv = meta.kv_lens_cuda.clone() + got = MTPDraftWorker._forward_small_batch_draft_loop( + worker, helper, inputs, meta, spec, object(), "draft", batch, counts + ) + selected = spec.batch_indices_cuda.long() * 6 + counts.long() - 1 + expected_tokens = (selected % 8).to(torch.int32).unsqueeze(1).expand(batch, 5) + torch.testing.assert_close(got, expected_tokens) + assert preflights == [("draft", batch * 6, batch)] + assert [call[1] for call in calls] == [batch * 6, batch, batch, batch, batch] + assert [call[4] for call in calls] == [False, True, True, True, True] + for step in range(1, 5): + torch.testing.assert_close(calls[step][2], original_positions[selected] + step) + torch.testing.assert_close(calls[step][3], original_kv - (5 - counts) + step - 1) + assert len(meta.updates) == 5 + for step, (seq_lens, kv_lens) in enumerate(meta.updates): + torch.testing.assert_close(seq_lens, torch.ones(batch, dtype=torch.int32)) + torch.testing.assert_close(kv_lens, original_kv - (5 - counts) + step) + assert meta.num_tokens == batch and not meta.use_spec_decoding + assert meta.kv_cache_manager == "target" + assert not meta.in_mtp_draft_loop and not meta.indexer_skip_topk + assert meta.mtp_num_accepted is None + assert worker._drafted_rows[0] is got + torch.testing.assert_close(worker._drafted_rows[1][:, 0], selected.to(torch.int32)) + torch.testing.assert_close(worker._drafted_rows[1][:, 1:], expected_tokens) + # The original inputs remain available to outer bookkeeping. + assert "owned_draft" not in inputs and "row_select" not in inputs + + +@pytest.mark.parametrize( + "batch,k,accepted", [(1, 3, [1]), (2, 3, [1, 4]), (1, 7, [1]), (1, 7, [8])] +) +def test_tp8_table_feedback_and_final_glue(reference_ops, batch, k, accepted): + worker, helper, inputs, meta, spec, counts, calls, preflights = _loop_fixture( + batch, tp_size=8, k=k, accepted=accepted + ) + original_kv = meta.kv_lens_cuda.clone() + got = MTPDraftWorker._forward_small_batch_draft_loop( + worker, helper, inputs, meta, spec, object(), "draft", batch, counts + ) + selected = spec.batch_indices_cuda.long() * (k + 1) + counts.long() - 1 + assert preflights == [("draft", batch * (k + 1), batch)] + assert [call[1] for call in calls] == [batch * (k + 1)] + [batch] * (k - 1) + assert [call[4] for call in calls] == [False] + [True] * (k - 1) + assert [(kind, step) for kind, step, _ in reference_ops] == [ + ("table", step) for step in range(k - 1) + ] + [("sample", k - 1)] + assert helper.projected_inputs[0] is None + torch.testing.assert_close(helper.head_inputs[0], inputs["hidden_states"][selected]) + for step in range(1, k): + prior = reference_ops[step - 1][2] + projected = helper.projected_inputs[step] + torch.testing.assert_close(projected[0], prior["normed_out"]) + torch.testing.assert_close(projected[1], prior["residual_out"]) + torch.testing.assert_close(helper.head_inputs[step], prior["normed_out"]) + torch.testing.assert_close(calls[step][2], inputs["position_ids"][selected] + step) + assert len(meta.updates) == k + for step, (seq_lens, kv_lens) in enumerate(meta.updates): + torch.testing.assert_close(seq_lens, torch.ones(batch, dtype=torch.int32)) + torch.testing.assert_close(kv_lens, original_kv - (k - counts) + step) + assert tuple(worker._drafted_rows[1].shape) == (batch, k + 1) + torch.testing.assert_close(worker._drafted_rows[1][:, 0], selected.to(torch.int32)) + torch.testing.assert_close(worker._drafted_rows[1][:, 1:], got) + assert meta.kv_cache_manager == "target" and not meta.use_spec_decoding + assert not meta.in_mtp_draft_loop and not meta.indexer_skip_topk + assert meta.mtp_num_accepted is None + + +def test_projected_output_bypasses_glue_gemm_and_input_hop(monkeypatch): + helper = object.__new__(small.SmallBatchDraftKernels) + helper.layer = object() + seen = [] + + def forbidden(*args, **kwargs): + raise AssertionError("projected rows must not execute the initial projection again") + + def attention(layer, positions, hidden, metadata, sharded, *, draft_state): + seen.append(hidden) + return hidden + 1 + + helper.decode = SimpleNamespace( + _hop=forbidden, + _attention=attention, + _post_attention_hop=lambda layer, hidden, residual: (hidden, residual), + ) + helper._moe_postnorm = lambda hidden, residual: hidden + residual + monkeypatch.setattr(small, "glm_draft_glue_rmsnorm_cat", forbidden) + monkeypatch.setattr(small, "low_m_gemm", forbidden) + normed, residual = torch.ones(2, 8), torch.full((2, 8), 3.0) + got = helper.forward_step( + None, None, torch.full_like(normed, -10), object(), eh_proj_output=(normed, residual) + ) + assert seen == [normed] + torch.testing.assert_close(got, normed + 1 + residual) + with pytest.raises(ValueError, match="mutually exclusive"): + helper.forward_step( + None, None, normed, object(), eh_proj_input=normed, eh_proj_output=(normed, residual) + ) + assert len(seen) == 1 + + +@pytest.mark.parametrize("rows", [1, 2]) +def test_tp8_head_launch_order_and_weighted_norm(monkeypatch, rows): + calls = [] + helper = object.__new__(small.SmallBatchDraftKernels) + helper.use_table = True + helper.hidden_rows = torch.ones(1, 8, dtype=torch.bfloat16) + helper.hnorm_weight = torch.full((8,), 2.0, dtype=torch.bfloat16) + helper.half_start = SimpleNamespace(record=lambda: calls.append("fork")) + helper.half_done = SimpleNamespace(record=lambda: calls.append("ready")) + aux = SimpleNamespace(wait_event=lambda event: calls.append("aux_wait")) + + class Norm: + variance_epsilon = 1e-5 + + def __call__(self, hidden): + calls.append("weighted_hnorm") + return hidden * helper.hnorm_weight + + helper.layer = SimpleNamespace(mlp=SimpleNamespace(aux_stream=aux), hnorm=Norm()) + helper.lm_head = object() + + def lm_head(lm, value): + calls.append("lm_head") + return value + + def half(value, weight, out, norm_weight, eps): + calls.append("hidden_half") + assert out.dtype == torch.float32 + assert (norm_weight is helper.hnorm_weight) == (rows == 1) + torch.testing.assert_close(value, torch.full_like(value, 1 if rows == 1 else 2)) + return out.fill_(7) + + helper.decode = SimpleNamespace(_lm_head_shard=lm_head) + monkeypatch.setattr(torch.cuda, "stream", lambda stream: contextlib.nullcontext()) + monkeypatch.setattr(small, "draft_hidden_half", half) + hidden = torch.ones(rows, 8, dtype=torch.bfloat16) + logits, projected = helper.head(hidden, next_step=True) + assert logits is hidden and tuple(projected.shape) == (rows, 1) + assert calls == ( + ["fork", "aux_wait", "hidden_half", "ready", "lm_head"] + if rows == 1 + else ["fork", "lm_head", "aux_wait", "weighted_hnorm", "hidden_half", "ready"] + ) + + +@pytest.mark.parametrize("tp_size,batch,k,fail_step", [(4, 3, 5, 2), (8, 2, 3, 2), (8, 1, 7, 6)]) +def test_exception_restores_draft_manager_and_index_share_without_replay( + reference_ops, tp_size, batch, k, fail_step +): + worker, helper, inputs, meta, spec, counts, calls, _ = _loop_fixture( + batch, fail_step=fail_step, tp_size=tp_size, k=k + ) + with pytest.raises(RuntimeError, match="injected draft failure"): + MTPDraftWorker._forward_small_batch_draft_loop( + worker, helper, inputs, meta, spec, object(), "draft", batch, counts + ) + assert [call[0] for call in calls] == list(range(fail_step + 1)) + assert len(meta.updates) == fail_step + assert meta.kv_cache_manager == "target" + assert not meta.in_mtp_draft_loop and not meta.indexer_skip_topk + assert meta.mtp_num_accepted is None and worker._drafted_rows is None + + +@pytest.mark.parametrize("tp_size,batch,k", [(4, 3, 5), (8, 2, 3), (8, 1, 7)]) +def test_rebound_manager_preflight_declines_before_any_forward(reference_ops, tp_size, batch, k): + worker, helper, inputs, meta, spec, counts, calls, preflights = _loop_fixture( + batch, eligible=False, tp_size=tp_size, k=k + ) + original_kv = meta.kv_lens_cuda.clone() + got = MTPDraftWorker._forward_small_batch_draft_loop( + worker, helper, inputs, meta, spec, object(), "draft", batch, counts + ) + assert got is None and calls == [] and meta.updates == [] + assert preflights == [("draft", batch * (spec.runtime_draft_len + 1), batch)] + torch.testing.assert_close(meta.kv_lens_cuda, original_kv) + assert meta.kv_cache_manager == "target" and not meta.in_mtp_draft_loop + assert meta.mtp_num_accepted is None + + +@pytest.mark.parametrize("batch,k", [(3, 5), (2, 3), (1, 3), (1, 7)]) +def test_first_attention_precedes_selection_of_hidden_and_residual( + reference_ops, monkeypatch, batch, k +): + helper = object.__new__(small.SmallBatchDraftKernels) + helper.hidden_size = 4 + helper.zero_residual = torch.zeros(24, 4) + helper.layer = SimpleNamespace( + eh_proj=SimpleNamespace(weight=torch.eye(4)), input_layernorm=object() + ) + seen = [] + + def attention(layer, positions, value, metadata, sharded, *, draft_state): + assert not sharded and draft_state is helper + seen.append(("attention", value.clone())) + return value + 10 + + helper.decode = SimpleNamespace( + _hop=lambda norm, partial, residual: (partial, partial + 100), + _attention=attention, + _post_attention_hop=lambda layer, output, residual: (output * 2, residual + 3), + ) + + def moe(value, residual): + seen.append(("moe", value.clone(), residual.clone())) + return value + + helper._moe_postnorm = moe + monkeypatch.setattr(small, "low_m_gemm", lambda x, weight, out: x @ weight.T) + full = torch.arange(batch * (k + 1) * 4, dtype=torch.float32).reshape(-1, 4) + order = torch.tensor([2, 0, 1] if batch == 3 else list(reversed(range(batch)))) + counts = torch.tensor([1, k + 1, 3][:batch]) + select = small.FirstDraftRowSelect(order, counts, k + 1) + got = helper.forward_step(None, None, full, object(), eh_proj_input=full, row_select=select) + expected_ids = order * (k + 1) + counts - 1 + torch.testing.assert_close(seen[0][1], full) + torch.testing.assert_close(seen[1][1], (full + 10)[expected_ids] * 2) + torch.testing.assert_close(seen[1][2], (full + 103)[expected_ids]) + torch.testing.assert_close(got, seen[1][1]) + with pytest.raises(RuntimeError, match="already selected"): + select(full, full) + + +def test_tp16_inline_dispatch_precedes_small_batch_branch(): + sentinel = object() + + def no_small_batch(*args): + raise AssertionError("the old inline path must win") + + worker = SimpleNamespace( + _drafting_target=lambda *args: sentinel, + _draft_on_target=lambda *args: "inline", + _small_batch_target=no_small_batch, + ) + result = MTPDraftWorker._forward_linear_draft_loop( + worker, {}, object(), object(), object(), None, 0, 1, object(), None + ) + assert result == ("inline", None) + + +@pytest.mark.parametrize("error_type", [MemoryError, OSError, RuntimeError]) +@pytest.mark.parametrize("tp_size", [4, 8]) +def test_local_draft_state_failure_reaches_group_vote(monkeypatch, error_type, tp_size): + comm = object() + votes = [] + decode = SimpleNamespace( + mapping=SimpleNamespace(tp_size=tp_size), + core=SimpleNamespace( + model_config=SimpleNamespace( + spec_config=SimpleNamespace(max_draft_len=3 if tp_size == 8 else 5) + ) + ), + ) + monkeypatch.setattr(small, "_get_mnnvl_workspace_comm", lambda mapping: comm) + monkeypatch.setattr(small, "_static_layout_supported", lambda *args: True) + + def vote(actual_comm, ok): + assert actual_comm is comm + votes.append(ok) + return ok + + def fail(*args, **kwargs): + raise error_type("local allocation failed") + + monkeypatch.setattr(small, "_all_succeeded", vote) + monkeypatch.setattr(small.DraftExchangeWorkspace, "for_sample_glue", lambda mapping: object()) + monkeypatch.setattr(small.DraftExchangeWorkspace, "for_ehproj_table", lambda mapping: object()) + monkeypatch.setattr(small, "SmallBatchDraftKernels", fail) + with pytest.raises(RuntimeError, match="failed on at least one rank") as caught: + small.build_small_batch_draft(decode, object(), object(), object()) + assert isinstance(caught.value.__cause__, error_type) + assert votes == [True, False] + + +def test_tp8_token_table_build_failure_reaches_group_vote(monkeypatch): + votes = [] + mapping = SimpleNamespace(tp_size=8, tp_rank=0) + norm = SimpleNamespace(weight=torch.ones(8, dtype=torch.bfloat16), variance_epsilon=1e-5) + decode = SimpleNamespace( + mapping=mapping, + core=SimpleNamespace( + embed_tokens=SimpleNamespace(weight=torch.ones(32, 8, dtype=torch.bfloat16)), + model_config=SimpleNamespace( + spec_config=SimpleNamespace(max_draft_len=3), + pretrained_config=SimpleNamespace( + hidden_size=8, num_experts_per_tok=8, routed_scaling_factor=2.5, index_topk=2 + ), + ), + ), + ) + layer = SimpleNamespace(enorm=norm, hnorm=norm) + weights = SimpleNamespace(w3_w1=object(), w2=object(), eh_proj=object()) + lm_head = SimpleNamespace(weight=torch.ones(4, 8, dtype=torch.bfloat16)) + monkeypatch.setattr(small, "_get_mnnvl_workspace_comm", lambda mapping: object()) + monkeypatch.setattr(small, "_static_layout_supported", lambda *args: True) + monkeypatch.setattr(small, "_all_succeeded", lambda comm, ok: votes.append(ok) or ok) + monkeypatch.setattr(small, "get_env_enable_pdl", lambda: True) + monkeypatch.setattr(small.DraftExchangeWorkspace, "for_sample_glue", lambda mapping: object()) + monkeypatch.setattr(small.DraftExchangeWorkspace, "for_ehproj_table", lambda mapping: object()) + + def fail_table(*args): + raise MemoryError("injected token-table allocation failure") + + monkeypatch.setattr(small, "build_draft_token_table", fail_table) + with pytest.raises(RuntimeError, match="failed on at least one rank") as caught: + small.build_small_batch_draft(decode, layer, lm_head, weights) + assert isinstance(caught.value.__cause__, MemoryError) + assert votes == [True, False] + + +@pytest.mark.parametrize( + "tp,batch,configured_k,runtime_k,helper_k,eligible", + [ + (4, 3, 5, 5, 5, True), + (8, 2, 3, 3, 3, True), + (8, 1, 7, 7, 7, True), + (8, 2, 7, 7, 7, False), + (8, 1, 5, 5, 5, False), + (4, 1, 7, 7, 7, False), + (8, 1, 7, 3, 7, False), + (8, 1, 7, 7, 3, False), + ], +) +def test_worker_checks_configured_runtime_and_allocated_draft_shape( + monkeypatch, tp, batch, configured_k, runtime_k, helper_k, eligible +): + # Tensor metadata stands in for device tensors; the real worker eligibility + # predicate runs before any kernel or collective can use these operands. + def operand(shape, dtype): + return SimpleNamespace( + shape=shape, + dtype=dtype, + device="cuda", + is_cuda=True, + is_contiguous=lambda: True, + data_ptr=lambda: 32, + dim=lambda: len(shape), + numel=lambda: shape[0], + ) + + rows = batch * (runtime_k + 1) + meta = _Metadata(batch, runtime_k) + monkeypatch.setattr(worker_module, "MTPIndexShareMetadata", _Metadata) + layer, lm_head = object(), object() + helper = SimpleNamespace( + layer=layer, + lm_head=lm_head, + DRAFT_LEN=helper_k, + MAX_BATCH=3 if tp == 4 else (1 if helper_k == 7 else 2), + ) + worker = SimpleNamespace( + mapping=SimpleNamespace(tp_size=tp), + max_draft_len=configured_k, + is_mtp_eagle=True, + spec_config=SimpleNamespace(draft_len_schedule=None), + guided_decoder=None, + sa_enhancer=None, + _d2t=None, + _decode_target=lambda metadata: SimpleNamespace(_small_batch_draft=helper), + ) + spec = SimpleNamespace( + runtime_draft_len=runtime_k, + batch_indices_cuda=operand((batch,), torch.int32), + wants_advanced_draft_sampling=False, + is_ragged_verify=False, + ) + inputs = dict( + input_ids=operand((rows,), torch.int32), + position_ids=operand((rows,), torch.int64), + hidden_states=operand((rows, 6144), torch.bfloat16), + ) + got = MTPDraftWorker._small_batch_target( + worker, + inputs, + meta, + spec, + SimpleNamespace(mtp_layers=[layer], lm_head=lm_head), + 0, + batch, + operand((batch,), torch.int32), + ) + assert got is (helper if eligible else None) + + +@pytest.mark.parametrize("configured_k", [None, 5, 8, [7]]) +def test_unsupported_tp8_draft_config_votes_before_workspace_allocation(monkeypatch, configured_k): + votes = [] + decode = SimpleNamespace( + mapping=SimpleNamespace(tp_size=8), + core=SimpleNamespace( + model_config=SimpleNamespace( + spec_config=None + if configured_k is None + else SimpleNamespace(max_draft_len=configured_k) + ) + ), + ) + monkeypatch.setattr(small, "_get_mnnvl_workspace_comm", lambda mapping: object()) + monkeypatch.setattr(small, "_all_succeeded", lambda comm, ok: votes.append(ok) or ok) + monkeypatch.setattr(small, "_static_layout_supported", lambda *args: True) + + def forbidden(*args, **kwargs): + raise AssertionError("unsupported draft configuration allocated a workspace") + + monkeypatch.setattr(small.DraftExchangeWorkspace, "for_sample_glue", forbidden) + monkeypatch.setattr(small.DraftExchangeWorkspace, "for_ehproj_table", forbidden) + assert small.build_small_batch_draft(decode, object(), object(), object()) is None + assert votes == [False] diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_verify_head_accept.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_verify_head_accept.py similarity index 98% rename from tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_verify_head_accept.py rename to tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_verify_head_accept.py index 9cef188ae5b6..96ccb07710d7 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_verify_head_accept.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_verify_head_accept.py @@ -9,7 +9,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.ops.spec.verify_head_accept import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.spec.verify_head_accept import ( # noqa: E501 glm_verify_head_accept as op, ) from tensorrt_llm._torch.speculative.interface import ( diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_verify_head_sample.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_verify_head_sample.py new file mode 100644 index 000000000000..c482c593c621 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/spec/test_modeling_v2_verify_head_sample.py @@ -0,0 +1,270 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""SM107 tokens-only FULL sampling: support, RNG identity, distributions and graph reuse. + +Run this file in separate processes with TRTLLM_ENABLE_PDL=0 and =1. Individual +sample IDs are compared only within this implementation, never across different +sampling algorithms. Dense logits retain the real154880-column stride throughout. +""" + +import math + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.spec.verify_head_sample import ( # noqa: E501 + VOCAB_SIZE, + VerifyHeadSampleState, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.ops.spec.verify_head_sample import ( # noqa: E501 + glm_verify_head_sample as op, +) +from tensorrt_llm._torch.pyexecutor.sampler.ops import custom as shared +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402,F401 + +pytestmark = pytest.mark.skipif(get_sm_version() != 107, reason="GLM Rubin sampling requires SM107") + + +def _inputs(rows, *, active=32, varied=True): + logits = torch.full((rows, VOCAB_SIZE), -1000.0, device="cuda") + values = ( + torch.linspace(0.0, -4.0, active, device="cuda") + if varied + else torch.zeros(active, device="cuda") + ) + logits[:, :active] = values + temps = torch.ones(rows, device="cuda") + ks = torch.full((rows,), torch.iinfo(torch.int32).max, device="cuda", dtype=torch.int32) + ps = torch.full((rows,), 0.95, device="cuda") + mins = torch.zeros(rows, device="cuda") + seeds = torch.full((rows,), 1337, device="cuda", dtype=torch.int64) + offsets = torch.arange(rows, device="cuda", dtype=torch.int64) * 32 + (1 << 62) + return [logits, temps, ks, ps, mins, seeds, offsets] + + +def _args(state, inputs): + return (*inputs, state.row_counters, state.workspace) + + +def _assert_support(tokens, probabilities): + assert tokens.dtype == torch.int32 and tokens.shape == (len(probabilities),) + assert ((tokens >= 0) & (tokens < VOCAB_SIZE)).all() + assert (probabilities.gather(1, tokens.long()[:, None]) > 0).all() + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + rows = cell.spec["rows"] + state = VerifyHeadSampleState("cuda") + inputs = _inputs(rows) + # Mixed rows drive split top-p, min-p only, top-k and the composed filtered path. + inputs[2][1::4] = 8 + inputs[4][2::4] = 0.15 + inputs[2][3::4] = 12 + inputs[4][3::4] = 0.05 + inputs[3][2::4] = 1 + certifying(op) + tokens = op(*_args(state, inputs)) + reference = op.reference(*_args(state, inputs)) + _assert_support(tokens, reference) + assert not state.row_counters.any() + + +@pytest.mark.parametrize("rows", [6, 12, 18]) +def test_rng_follows_rows_across_permutation_and_split(rows): + state = VerifyHeadSampleState("cuda") + inputs = _inputs(rows) + inputs[2][1::3] = 7 + inputs[4][2::3] = 0.1 + whole = state.sample(*inputs) + again = state.sample(*inputs) + assert torch.equal(whole, again) + order = torch.arange(rows - 1, -1, -1, device="cuda") + permuted = state.sample(*(x[order].contiguous() for x in inputs)) + assert torch.equal(permuted, whole[order]) + # Singleton must use the per-row stream, including offsets in the unseeded domain. + singles = torch.cat([state.sample(*(x[i : i + 1] for x in inputs)) for i in range(rows)]) + assert torch.equal(singles, whole) + parts = torch.cat( + [state.sample(*(x[a:b] for x in inputs)) for a, b in ((0, rows // 2), (rows // 2, rows))] + ) + assert torch.equal(parts, whole) + assert not state.row_counters.any() + + +def test_scalar_rng_separates_rows_and_per_row_rng_does_not(): + state = VerifyHeadSampleState("cuda") + inputs = _inputs(24, varied=False) + inputs[3].fill_(1) + inputs[6].fill_(64) + same_stream = state.sample(*inputs) + assert torch.equal(same_stream, same_stream[:1].expand(24)) + scalar_inputs = [*inputs[:5], inputs[5][:1], inputs[6][:1]] + scalar = state.sample(*scalar_inputs) + assert torch.equal(scalar, state.sample(*scalar_inputs)) + assert scalar.unique().numel() > 1 + # Per-token offsets change subsequences rather than overlapping draws within one stream. + inputs[6].copy_(torch.arange(24, device="cuda", dtype=torch.int64)) + assert state.sample(*inputs).unique().numel() > 1 + + +@pytest.mark.parametrize( + "k,p,min_p", [(0, 0.8, 0.0), (8, 0.8, 0.05), (0, 1.0, 0.15), (7, 1.0, 0.0)] +) +def test_distribution_on_real_verify_rows(k, p, min_p): + state = VerifyHeadSampleState("cuda") + inputs = _inputs(18, active=16) + inputs[2].fill_(k) + inputs[3].fill_(p) + inputs[4].fill_(min_p) + expected = op.reference(*_args(state, inputs))[0].cpu() + samples = [] + for iteration in range(128): + inputs[6].copy_(torch.arange(18, device="cuda", dtype=torch.int64) + iteration * 18) + samples.append(state.sample(*inputs)) + tokens = torch.cat(samples) + counts = torch.bincount(tokens.long(), minlength=VOCAB_SIZE).float().cpu() + assert not counts[expected == 0].any() + measured = counts / tokens.numel() + # 2304 independent subsequences: TV0.08 is wider than the expected multinomial + # fluctuation for <=16 bins, yet rejects missing filters or deterministic fallback. + assert (measured - expected).abs().sum().item() * 0.5 < 0.08 + assert not state.row_counters.any() + + +@pytest.mark.parametrize("direction", [-1, 0, 1]) +def test_top_p_boundary_keeps_complete_ties(direction): + state = VerifyHeadSampleState("cuda") + inputs = _inputs(24, active=4, varied=False) + inputs[0][:, 2:4] = -math.log(2.0) + boundary = torch.tensor(2.0 / 3.0, dtype=torch.float32, device="cuda") + if direction: + boundary = torch.nextafter(boundary, torch.tensor(float("inf") * direction, device="cuda")) + inputs[3].fill_(boundary) + # Above each half-weight token is exactly two unit weights. FULL admits the + # whole half-weight tie at mass==top_p*total; one ULP below excludes both. + samples = [] + for i in range(16): + inputs[6].add_(24) + samples.append(state.sample(*inputs)) + tokens = torch.cat(samples) + if direction < 0: + assert (tokens < 2).all() + else: + assert (tokens < 4).all() and (tokens == 2).any() and (tokens == 3).any() + # The current-main materialized FULL support is an independent implementation. + current = shared.fused_compute_probs_from_logits(*inputs[:5]) + assert (current[:, :2] > 0).all() + assert bool((current[:, 2:4] > 0).all()) == (direction >= 0) + assert not state.row_counters.any() + + +@pytest.mark.parametrize("min_p", [0.0, 1.0, 2.0]) +def test_top_k_and_min_p_retain_weight_ties(min_p): + state = VerifyHeadSampleState("cuda") + inputs = _inputs(24, active=8, varied=False) + inputs[0][:, 4:8] = -1 + inputs[2].fill_(1) + inputs[3].fill_(0.0) + inputs[4].fill_(min_p) + out = state.sample(*inputs) + assert (out < 4).all() and out.unique().numel() > 1 + assert not state.row_counters.any() + + +def test_graph_replay_changes_filter_roles_without_recapture(): + state = VerifyHeadSampleState("cuda") + inputs = _inputs(18) + # Warm-up on a side stream as required by CUDA graph capture. + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + state.sample(*inputs) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + output = state.sample(*inputs) + for min_p in (0.0, 0.05, 1.0, 0.0): + inputs[4].fill_(min_p) + inputs[2].fill_(0) + inputs[2][::3] = 8 if min_p != 1 else 1 + inputs[3].fill_(0.95) + inputs[3][1::3] = 1 + inputs[1].fill_(0.7 if min_p else 1.0) + inputs[6].add_(32) + # Sampling params are completed step inputs and can change between replays. + graph.replay() + actual = output.clone() + expected = state.sample(*inputs) + assert torch.equal(actual, expected) + _assert_support(actual, op.reference(*_args(state, inputs))) + assert not state.row_counters.any() + assert state.workspace.numel() == 1440384 + + +def test_delayed_logits_producer_and_persistent_counter_reuse(): + state = VerifyHeadSampleState("cuda") + inputs = _inputs(18) + original = inputs[0].clone() + baseline = state.sample(*inputs) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + inputs[0].fill_(-1000) + torch.cuda._sleep(1000000) + inputs[0].copy_(original) + result = state.sample(*inputs) + torch.cuda.current_stream().wait_stream(stream) + assert torch.equal(result, baseline) + # Reusing one max-capacity arena with different row counts must not preserve stale + # counters or depend on per-call scratch offsets from the previous shape. + for rows in (1, 24, 6, 18, 12, 1): + batch = _inputs(rows) + actual = state.sample(*batch) + assert torch.equal(actual, state.sample(*batch)) + assert not state.row_counters.any() + + +def test_input_contract_and_mutation_schema(): + state = VerifyHeadSampleState("cuda") + inputs = _inputs(6) + assert state.supported(*inputs) + bad = list(inputs) + bad[2] = bad[2].long() + assert not state.supported(*bad) + bad = list(inputs) + bad[5] = bad[5][:2] + assert not state.supported(*bad) + bad = list(inputs) + bad[0] = bad[0][:, ::2] + assert not state.supported(*bad) + with pytest.raises(ValueError): + state.sample(*bad) + schema = torch.ops.trtllm.glm_verify_head_sample_tokens.default._schema + writes = {a.name for a in schema.arguments if a.alias_info and a.alias_info.is_write} + assert writes == {"row_counters", "workspace"} + with pytest.raises(RuntimeError, match="workspace"): + op(*inputs, state.row_counters, state.workspace[:16]) + + +@pytest.mark.parametrize("k,min_p", [(0, 0.0), (1, 0.0), (8, 0.05), (VOCAB_SIZE, 1.0)]) +def test_masked_logits_have_zero_weight(k, min_p, certifying): # noqa: F811 + state = VerifyHeadSampleState("cuda") + inputs = _inputs(18, active=4) + inputs[0][:, 4:] = -float("inf") + # Masked holes and a singleton finite row exercise empty partial segments. + inputs[0][::2, 1] = -float("inf") + inputs[0][-1].fill_(-float("inf")) + inputs[0][-1, VOCAB_SIZE - 1] = 2.0 + inputs[2].fill_(k) + inputs[4].fill_(min_p) + certifying(op) + tokens = op(*_args(state, inputs)) + _assert_support(tokens, op.reference(*_args(state, inputs))) + assert tokens[-1] == VOCAB_SIZE - 1 + assert torch.isfinite(inputs[0].gather(1, tokens.long()[:, None])).all() + assert not state.row_counters.any() diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/test_routing.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/test_routing.py index 5e83f72e3020..0b3ccfc81387 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/test_routing.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/test_routing.py @@ -59,7 +59,8 @@ def _model_config( _GLM52_TARGET = ( - "models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__tp16.modeling.ModelingV2Glm52Nvfp4Sm107Tp16" + "models.glm_moe_dsa.glm_5_2_nvfp4__sm_107__low_latency.modeling." + "ModelingV2Glm52Nvfp4Sm107LowLatency" ) _TP16 = dict(world_size=16, tp_size=16, moe_ep_size=1, moe_tp_size=16) @@ -84,15 +85,23 @@ def _glm52_config(**overrides): return PretrainedConfig(**fields) +@pytest.mark.parametrize("tp", [4, 8, 16]) @pytest.mark.parametrize("mode", ["auto", "require"]) @pytest.mark.parametrize( "spec_config", [None, MTPDecodingConfig(max_draft_len=7), MTPDecodingConfig(max_draft_len=3)], ids=["no-spec", "mtp7", "mtp3"], ) -def test_glm52_tp16_matches(_on_sm107, spec_config, mode): +def test_glm52_supported_tp_matches(_on_sm107, spec_config, mode, tp): """Any draft length: the draft loop replays the one MTP module per step.""" - config = _model_config(_glm52_config(), spec_config=spec_config, **_TP16) + config = _model_config( + _glm52_config(), + spec_config=spec_config, + world_size=tp, + tp_size=tp, + moe_ep_size=1, + moe_tp_size=tp, + ) assert _qualified(modeling_v2_resolve(config, _llm_args(mode))) == _GLM52_TARGET @@ -102,11 +111,11 @@ def test_glm52_tp16_matches(_on_sm107, spec_config, mode): (dict(num_hidden_layers=92), _TP16, None, r"shape \(92, 6144, 256, 2048\)"), (dict(), dict(_TP16, enable_attention_dp=True), None, "attention_dp=True"), (dict(), dict(_TP16, moe_ep_size=16, moe_tp_size=1), None, "moe_ep=16"), - (dict(), dict(world_size=8, tp_size=8, moe_tp_size=8), None, "world_size=8"), + (dict(), dict(world_size=2, tp_size=2, moe_tp_size=2), None, "world_size=2"), (dict(), _TP16, MTPDecodingConfig(max_draft_len=3, use_mtp_vanilla=True), "MTP"), (dict(), _TP16, _NGRAM, "NGRAM"), ], - ids=["other-shape", "attention-dp", "expert-parallel", "tp8", "vanilla-mtp", "ngram"], + ids=["other-shape", "attention-dp", "expert-parallel", "tp2", "vanilla-mtp", "ngram"], ) def test_glm52_near_misses_do_not_match( _on_sm107, config_kwargs, mapping_kwargs, spec_config, missed diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py index 067903a2ba20..5330cdc0d90f 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py @@ -427,7 +427,7 @@ def test_every_target_has_the_accuracy_gate_of_its_name(): # `unittest/`, with what they wait for. They are the only tests no list may # reach, and none of them may be listed. _NO_CI_STAGE = { - "unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm": ( + "unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__low_latency/comm": ( "4 ranks on sm_107; CI has only the single-GPU GR100 stage (auto:gr100-x1)" ), }